Update app.py

app.py

@@ -5,40 +5,43 @@ from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
 
-num_classes = 2  
+num_classes = 2  # Number of classes for your dataset
 
+# Download model weights from Hugging Face
 def download_model():
     model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
     return model_path
 
+# Load the model from the downloaded weights
 def load_model(model_path):
-    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() 
+    model = models.resnet50(pretrained=False)  # Set pretrained=False for custom weights
+    model.fc = nn.Linear(model.fc.in_features, num_classes)  # Adjust final layer for your number of classes
+    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))  # Load model weights
+    model.eval()  # Set model to evaluation mode
     return model
 
-model_path = download_model()  
+# Download and load the model
+model_path = download_model()
 model = load_model(model_path)
 
-
+# Image transformation pipeline
 transform = transforms.Compose([
     transforms.Resize(256),  # Resize the image to 256x256
     transforms.CenterCrop(224),  # Crop the image to 224x224
     transforms.ToTensor(),  # Convert the image to a Tensor
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # Normalize for ImageNet
 ])
 
+# Prediction function
 def predict(image):
-    
-    image = transform(image).unsqueeze(0)  
-    image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))  
-    
+    image = transform(image).unsqueeze(0)  # Add batch dimension
+    image = image.to(torch.device("cpu"))  # Move the image to CPU (adjust if you want to use GPU)
+
     with torch.no_grad():
-        outputs = model(image)  
-        predicted_class = torch.argmax(outputs, dim=1).item()  
+        outputs = model(image)  # Perform forward pass
+        predicted_class = torch.argmax(outputs, dim=1).item()  # Get the predicted class ID
 
-    
+    # Return appropriate response based on predicted class
     if predicted_class == 0:
         return "The photo you've sent is of fall army worm with problem ID 126."
     elif predicted_class == 1:
@@ -46,13 +49,16 @@ def predict(image):
     else:
         return "Unexpected class prediction."
 
+# Create the Gradio interface and expose it as an API
 iface = gr.Interface(
-    fn=predict,  
-    inputs=gr.Image(type="pil"),  
-    outputs=gr.Textbox(),  
-    live=True,  
+    fn=predict,  # Prediction function
+    inputs=gr.Image(type="pil"),  # Image input (PIL format)
+    outputs=gr.Textbox(),  # Text output (Predicted class description)
+    live=True,  # Update predictions as the user uploads an image
     title="Maize Anomaly Detection",
-    description="Upload an image of maize to detect anomalies like disease or pest infestation."
+    description="Upload an image of maize to detect anomalies like disease or pest infestation.",
+    api=True  # Expose the Gradio interface for API calls (POST requests)
 )
 
+# Launch the Gradio interface
 iface.launch()