Update app.py

app.py

@@ -5,43 +5,45 @@ from torchvision import models, transforms
 from huggingface_hub import hf_hub_download
 from PIL import Image
 
-num_classes = 2  # Number of classes for your dataset
+# Define the number of classes
+num_classes = 2  # Update with the actual number of classes in your dataset (e.g., 2 for healthy and anomalous)
 
-# Download model weights from Hugging Face
+# Download model from Hugging Face
 def download_model():
-    model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
+    model_path = hf_hub_download(repo_id="your_huggingface_username/your_model_name", filename="pytorch_model.bin")
     return model_path
 
-# Load the model from the downloaded weights
+# Load the model from Hugging Face
 def load_model(model_path):
-    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
+    model = models.resnet50(pretrained=False)  # Set pretrained=False because you're loading custom weights
+    model.fc = nn.Linear(model.fc.in_features, num_classes)  # Adjust for the number of classes in your dataset
+    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))  # Load model on CPU for compatibility
+    model.eval()  # Set to evaluation mode
     return model
 
-# Download and load the model
-model_path = download_model()
+# Download the model and load it
+model_path = download_model()  # Downloads the model from Hugging Face Hub
 model = load_model(model_path)
 
-# Image transformation pipeline
+# Define the transformation for the input image
 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]),  # Normalize for ImageNet
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # Normalize the image (ImageNet mean and std)
 ])
 
-# Prediction function
+# Define the prediction function
 def predict(image):
+    # Apply the necessary transformations to the image
     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)
-
+    image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))  # Move to GPU if available
+    
     with torch.no_grad():
         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
+        predicted_class = torch.argmax(outputs, dim=1).item()  # Get the predicted class
+    
+    # Create a response based on the 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:
@@ -51,14 +53,13 @@ def predict(image):
 
 # Create the Gradio interface
 iface = gr.Interface(
-    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."
+    fn=predict,  # Function for prediction
+    inputs=gr.Image(type="pil"),  # Image input
+    outputs=gr.Textbox(),  # Output: Predicted class
+    live=True,  # Updates as the user uploads an image
+    title="Wheat Anomaly Detection",
+    description="Upload an image of wheat to detect anomalies like disease or pest infestation."
 )
 
-# Expose Gradio interface as API endpoint
-iface.launch
-
+# Launch the Gradio interface
+iface.launch()