Update app.py

app.py

@@ -1,80 +1,32 @@
-import gradio as gr
-import json
-import torch
-from torch import nn
-from torchvision import models, transforms
-from huggingface_hub import hf_hub_download
-from PIL import Image
-import requests
-import base64
-from io import BytesIO
-import os
-
-# Define the number of classes
-num_classes = 2
-
-# Download model from Hugging Face
-def download_model():
-    try:
-        model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
-        return model_path
-    except Exception as e:
-        print(f"Error downloading model: {e}")
-        return None
-
-# Load the model from Hugging Face
-def load_model(model_path):
+def predict(data):
     try:
-        model = models.resnet50(pretrained=False)
-        model.fc = nn.Linear(model.fc.in_features, num_classes)
-        model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
-        model.eval()
-        return model
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None
-
-# Download the model and load it
-model_path = download_model()
-model = load_model(model_path) if model_path else None
+        image_input = data.get('image', None)
+        if not image_input:
+            return json.dumps({"error": "No image provided."})
 
-# Define the transformation for the input image
-transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-])
-
-def predict(image):
-    try:
-        print(f"Received image input: {image}")
+        print(f"Received image input: {image_input}")
 
         # Check if the input is a PIL Image type
-        if isinstance(image, Image.Image):
-            print(f"Image is already loaded as PIL Image: {image}")
+        if isinstance(image_input, Image.Image):
+            print(f"Image is already loaded as PIL Image: {image_input}")
         else:
             # Check if the input contains a base64-encoded string or URL
-            if isinstance(image, dict) and image.get("data"):
-                image_data = image["data"]
-                if image_data.startswith("http"):  # URL case
-                    response = requests.get(image_data)
+            if image_input.startswith("http"):  # URL case
+                try:
+                    response = requests.get(image_input)
                     image = Image.open(BytesIO(response.content))
                     print(f"Fetched image from URL: {image}")
-                else:  # Base64-encoded image case
-                    image_data = base64.b64decode(image_data)
+                except Exception as e:
+                    print(f"Error fetching image from URL: {e}")
+                    return json.dumps({"error": f"Error fetching image from URL: {e}"})
+            else:  # Assuming it is base64-encoded image data
+                try:
+                    image_data = base64.b64decode(image_input)
                     image = Image.open(BytesIO(image_data))
                     print(f"Decoded base64 image: {image}")
-            elif isinstance(image, str) and image.startswith("http"):
-                response = requests.get(image)
-                image = Image.open(BytesIO(response.content))
-                print(f"Fetched image from URL: {image}")
-            elif isinstance(image, str) and os.path.isfile(image):
-                image = Image.open(image)
-                print(f"Loaded image from local path: {image}")
-            else:
-                print("Invalid image format received.")
-                return json.dumps({"error": "Invalid image format received."})
+                except Exception as e:
+                    print(f"Error decoding base64 image: {e}")
+                    return json.dumps({"error": f"Error decoding base64 image: {e}"})
 
         # Apply transformations
         image = transform(image).unsqueeze(0)
@@ -109,4 +61,3 @@ iface = gr.Interface(
 
 # Launch the Gradio interface 
 iface.launch(share=True, show_error=True)
-