Update app.py

app.py

@@ -2,6 +2,7 @@ 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
@@ -11,19 +12,30 @@ import os
 # Define the number of classes
 num_classes = 2  # Update with the actual number of classes in your dataset
 
-# Load the model (assuming you've already downloaded it)
-def load_model():
+# 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):
     try:
         model = models.resnet50(pretrained=False)
         model.fc = nn.Linear(model.fc.in_features, num_classes)
-        model.load_state_dict(torch.load("path_to_your_model.pth", map_location=torch.device("cpu")))
+        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
 
-model = load_model()
+# Download the model and load it
+model_path = download_model()
+model = load_model(model_path) if model_path else None
 
 # Define the transformation for the input image
 transform = transforms.Compose([
@@ -33,41 +45,54 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-# Prediction function
-def process_image(image, image_url=None):
+def predict(image):
     try:
-        # Ensure that the image is not None
-        if image is None and not image_url:
-            return "No image or URL provided."
+        print(f"Received image input: {image}")
 
-        # Handle URL-based image loading
-        if image_url:
+        # Check if the input contains a base64-encoded string
+        if isinstance(image, dict) and image.get("data"):
             try:
-                response = requests.get(image_url)
-                response.raise_for_status()  # Raise an error if the request fails
+                image_data = base64.b64decode(image["data"])
+                image = Image.open(BytesIO(image_data))
+                print(f"Decoded base64 image: {image}")
+            except Exception as e:
+                print(f"Error decoding base64 image: {e}")
+                return f"Error decoding base64 image: {e}"
+
+        # Check if the input is a URL
+        elif isinstance(image, str) and image.startswith("http"):
+            try:
+                response = requests.get(image)
                 image = Image.open(BytesIO(response.content))
+                print(f"Fetched image from URL: {image}")
             except Exception as e:
+                print(f"Error fetching image from URL: {e}")
                 return f"Error fetching image from URL: {e}"
 
-        # Handle local file path image loading (Gradio File input)
+        # Check if the input is a local file path
         elif isinstance(image, str) and os.path.isfile(image):
             try:
                 image = Image.open(image)
+                print(f"Loaded image from local path: {image}")
             except Exception as e:
+                print(f"Error loading image from local path: {e}")
                 return f"Error loading image from local path: {e}"
 
-        # Validate that the image is loaded correctly
+        # Validate that the image is correctly loaded
         if not isinstance(image, Image.Image):
+            print("Invalid image format received.")
             return "Invalid image format received."
 
         # Apply transformations
         image = transform(image).unsqueeze(0)
+        print(f"Transformed image tensor: {image.shape}")
 
-        # 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()
+            print(f"Prediction output: {outputs}, Predicted class: {predicted_class}")
 
         if predicted_class == 0:
             return "The photo you've sent is of fall army worm with problem ID 126."
@@ -76,20 +101,18 @@ def process_image(image, image_url=None):
         else:
             return "Unexpected class prediction."
     except Exception as e:
+        print(f"Error processing image: {e}")
         return f"Error processing image: {e}"
 
 # Create the Gradio interface
 iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.File(label="Upload an image (Local File Path)"),  # Input: Local file
-        gr.Textbox(label="Enter Image URL", placeholder="Enter image URL here", lines=1)  # Input: Image URL
-    ],
-    outputs=gr.Textbox(label="Prediction Result"),  # Output: Prediction result
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload an image or provide a URL or local path"),  # Input: Image, URL, or Local Path
+    outputs=gr.Textbox(label="Prediction Result"),  # Output: Predicted class
     live=True,
     title="Maize Anomaly Detection",
-    description="Upload an image of maize to detect anomalies like disease or pest infestation. You can upload local images or provide an image URL."
+    description="Upload an image of maize to detect anomalies like disease or pest infestation. You can provide local paths, URLs, or base64-encoded images."
 )
 
 # Launch the Gradio interface
-iface.launch(share=True, show_error=True)
+iface.launch(share=True, show_error=True)