Update app.py

app.py

@@ -7,7 +7,6 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 import requests
 import os
-import base64
 from io import BytesIO
 
 # Define the number of classes
@@ -15,28 +14,20 @@ num_classes = 2
 
 # 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
+    model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
+    return model_path
 
 # 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(model_path, map_location=torch.device("cpu")))
-        model.eval()
-        return model
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None
+    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()
+    return model
 
 # Download the model and load it
 model_path = download_model()
-model = load_model(model_path) if model_path else None
+model = load_model(model_path)
 
 # Define the transformation for the input image
 transform = transforms.Compose([
@@ -46,48 +37,50 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-def predict(image):
-    try:
-        print(f"Received image input: {image}")
-
-        # Check if the input is a PIL Image type (Gradio automatically provides a PIL image)
-        if not isinstance(image, Image.Image):
-            return json.dumps({"error": "Invalid image format received. Please provide a valid image."})
+# Function to predict from image content
+def predict_from_image(image):
+    # Ensure the image is a PIL Image
+    if not isinstance(image, Image.Image):
+        raise ValueError("Invalid image format received. Please provide a valid image.")
 
-        # Apply transformations to the image
-        image = transform(image).unsqueeze(0)
-        print(f"Transformed image tensor: {image.shape}")
+    # Apply transformations
+    image_tensor = transform(image).unsqueeze(0)
 
-        # Move the image to the correct device
-        image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    # Predict
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        predicted_class = torch.argmax(outputs, dim=1).item()
 
-        # Make predictions
-        with torch.no_grad():
-            outputs = model(image)
-            predicted_class = torch.argmax(outputs, dim=1).item()
-            print(f"Prediction output: {outputs}, Predicted class: {predicted_class}")
-
-        # Return the result based on the predicted class
-        if predicted_class == 0:
-            return json.dumps({"result": "The photo you've sent is of fall army worm with problem ID 126."})
-        elif predicted_class == 1:
-            return json.dumps({"result": "The photo you've sent is of a healthy maize image."})
-        else:
-            return json.dumps({"error": "Unexpected class prediction."})
+    # 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 is of a healthy maize image."}
+    else:
+        return {"error": "Unexpected class prediction."}
 
+# Function to predict from URL
+def predict_from_url(url):
+    try:
+        response = requests.get(url)
+        response.raise_for_status()  # Ensure the request was successful
+        image = Image.open(BytesIO(response.content))
+        return predict_from_image(image)
     except Exception as e:
-        print(f"Error processing image: {e}")
-        return json.dumps({"error": f"Error processing image: {e}"})
+        return {"error": f"Failed to process the URL: {str(e)}"}
 
-# Create the Gradio interface
+# Gradio interface
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an image or provide a URL or local path"),
-    outputs=gr.Textbox(label="Prediction Result"),
+    fn=lambda image, url: predict_from_image(image) if image else predict_from_url(url),
+    inputs=[
+        gr.Image(type="pil", label="Upload an Image"),
+        gr.Textbox(label="Or Enter an Image URL", placeholder="Provide a valid image URL"),
+    ],
+    outputs=gr.JSON(label="Prediction Result"),
     live=True,
     title="Maize Anomaly Detection",
-    description="Upload an image of maize to detect anomalies like disease or pest infestation. You can provide local paths, URLs, or base64-encoded images."
+    description="Upload an image or provide a URL to detect anomalies in maize crops.",
 )
 
-# Launch the Gradio interface 
+# Launch the interface
 iface.launch(share=True, show_error=True)