Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 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
@@ -12,30 +11,19 @@ import os
 # Define the number of classes
 num_classes = 2  # Update with the actual number of classes in your dataset
 
-# 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):
+# Load the model (assuming you've already downloaded it)
+def load_model():
     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.load_state_dict(torch.load("path_to_your_model.pth", 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
+model = load_model()
 
 # Define the transformation for the input image
 transform = transforms.Compose([
@@ -45,46 +33,45 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-def process_image(image_input):
+# Prediction function
+def process_image(data):
     try:
-        # Process the image input (URL, local file, or base64)
-        if isinstance(image_input, dict):
-            # Check if the input contains a URL
-            if image_input.get("url"):
-                image_url = image_input["url"]
-                response = requests.get(image_url)
-                image = Image.open(BytesIO(response.content))
-            # Check if the input contains a file path
-            elif image_input.get("path"):
-                image_path = image_input["path"]
-                image = Image.open(image_path)
-            # Handle base64 if it's included
-            elif image_input.get("data"):
-                image_data = base64.b64decode(image_input["data"])
+        # Check if the input contains a base64-encoded string
+        if isinstance(data, dict):
+            if "data" in data:
+                # Base64 decoding
+                image_data = base64.b64decode(data["data"])
                 image = Image.open(BytesIO(image_data))
+            elif "url" in data:
+                # URL-based image loading
+                response = requests.get(data["url"])
+                image = Image.open(BytesIO(response.content))
+            elif "path" in data:
+                # Local path image loading
+                image = Image.open(data["path"])
             else:
-                return "Invalid input data format. Please provide a URL or path."
+                return "Invalid input data structure."
 
-            # Apply transformations
-            image = transform(image).unsqueeze(0)
-            image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        # Validate image
+        if not isinstance(image, Image.Image):
+            return "Invalid image format received."
 
-            # Make the prediction
-            with torch.no_grad():
-                outputs = model(image)
-                predicted_class = torch.argmax(outputs, dim=1).item()
+        # Apply transformations
+        image = transform(image).unsqueeze(0)
 
-            # Return prediction result
-            if predicted_class == 0:
-                return "The photo you've sent is of fall army worm with problem ID 126."
-            elif predicted_class == 1:
-                return "The photo you've sent is of a healthy maize image."
-            else:
-                return "Unexpected class prediction."
+        # Prediction
+        image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        with torch.no_grad():
+            outputs = model(image)
+            predicted_class = torch.argmax(outputs, dim=1).item()
+
+        if predicted_class == 0:
+            return "The photo you've sent is of fall army worm with problem ID 126."
+        elif predicted_class == 1:
+            return "The photo you've sent is of a healthy maize image."
         else:
-            return "Invalid input. Please provide a dictionary with 'url' or 'path'."
+            return "Unexpected class prediction."
     except Exception as e:
-        print(f"Error processing image: {e}")
         return f"Error processing image: {e}"
 
 # Create the Gradio interface