Update app.py

app.py

@@ -4,6 +4,7 @@ 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
@@ -15,18 +16,11 @@ logging.basicConfig(level=logging.INFO)
 num_classes = 3
 
 # Download model from Hugging Face
-
 def download_model():
-    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
+    model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
+    return model_path
 
 # Load the model from Hugging Face
-
 def load_model(model_path):
     model = models.resnet50(pretrained=False)
     num_features = model.fc.in_features
@@ -34,7 +28,7 @@ def load_model(model_path):
         nn.Dropout(0.5),
         nn.Linear(num_features, 3)  # 3 classes
     )
-    
+
     # Load the checkpoint
     checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
     
@@ -50,9 +44,6 @@ def load_model(model_path):
     model.load_state_dict(new_state_dict, strict=False)
     model.eval()
     return model
-    except Exception as e:
-        logging.error(f"Failed to load model: {e}")
-        raise
 
 # Path to your model
 model_path = download_model()
@@ -73,7 +64,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)).convert('RGB')
+        image = Image.open(BytesIO(response.content))
 
         # Apply transformations
         image_tensor = transform(image).unsqueeze(0)
@@ -81,24 +72,22 @@ 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
-        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.")}
+        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."}
 
     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",
@@ -108,4 +97,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()