Update app.py

app.py

@@ -4,7 +4,6 @@ from torch import nn
 from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
-import os
 import logging
 import requests
 from io import BytesIO
@@ -16,31 +15,43 @@ logging.basicConfig(level=logging.INFO)
 num_classes = 3
 
 # Download model from Hugging Face
+
 def download_model():
-    model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
-    return model_path
+    try:
+        model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
+        logging.info("Model downloaded successfully.")
+        return model_path
+    except Exception as e:
+        logging.error(f"Failed to download model: {e}")
+        raise
 
 # Load the model from Hugging Face
-def load_model(model_path):
-    model = models.resnet50(pretrained=False)
-    num_features = model.fc.in_features
-    model.fc = nn.Sequential(
-        nn.Dropout(0.5),
-        nn.Linear(num_features, 3)  # 3 classes
-    )
-    checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
-    model.load_state_dict(checkpoint['model_state_dict'])
-    
-    model.eval()
-    return model
 
+def load_model(model_path):
+    try:
+        model = models.resnet50(pretrained=False)
+        num_features = model.fc.in_features
+        model.fc = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(num_features, num_classes)
+        )
+        checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
+
+        # Ensure compatibility by handling key mismatches
+        model_state_dict = checkpoint['model_state_dict']
+        for key in list(model_state_dict.keys()):
+            if key.startswith('fc.1'):
+                model_state_dict[key.replace('fc.1', 'fc')] = model_state_dict.pop(key)
+
+        model.load_state_dict(model_state_dict)
+        model.eval()
+        logging.info("Model loaded successfully.")
+        return model
+    except Exception as e:
+        logging.error(f"Failed to load model: {e}")
+        raise
 
 # Path to your model
-model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
-model = load_model(model_path)
-
-
-# Download the model and load it
 model_path = download_model()
 model = load_model(model_path)
 
@@ -59,7 +70,7 @@ def predict_from_image_url(image_url):
         # Download the image from the provided URL
         response = requests.get(image_url)
         response.raise_for_status()
-        image = Image.open(BytesIO(response.content))
+        image = Image.open(BytesIO(response.content)).convert('RGB')
 
         # Apply transformations
         image_tensor = transform(image).unsqueeze(0)
@@ -67,27 +78,29 @@ def predict_from_image_url(image_url):
         # Perform prediction
         with torch.no_grad():
             outputs = model(image_tensor)
+            if outputs.shape[1] != num_classes:
+                raise ValueError(f"Unexpected number of output classes: {outputs.shape[1]} (expected {num_classes})")
             predicted_class = torch.argmax(outputs, dim=1).item()
 
         # Interpret the result
-        if predicted_class == 0:
-            return {"result": "The photo is of Fall Army Worm with problem ID 126."}
-        elif predicted_class == 1:
-            return {"result": "The photo shows symptoms of Phosphorus Deficiency with Problem ID 142."}
-        elif predicted_class == 2:
-            return {"result": "The photo shows symptoms of Bacterial Leaf Blight with Problem ID 203."}
-        else:
-            return {"error": "Unexpected class prediction."}
+        class_map = {
+            0: "The photo is of Fall Army Worm with problem ID 126.",
+            1: "The photo shows symptoms of Phosphorus Deficiency with Problem ID 142.",
+            2: "The photo shows symptoms of Bacterial Leaf Blight with Problem ID 203."
+        }
+        return {"result": class_map.get(predicted_class, "Unexpected class prediction.")}
 
     except Exception as e:
+        logging.error(f"Error during prediction: {e}")
         return {"error": str(e)}
 
 
+# Initialize Gradio interface
 demo = gr.Interface(
     fn=predict_from_image_url,
     inputs="text",
     outputs="json",
-    title="Maize Disease Classification",
+    title="Crop anomaly Classification",
     description="Enter a URL to an image for classification (Fall Army Worm, Phosphorus Deficiency, or Bacterial Leaf Blight).",
 )