Update app.py

app.py

@@ -36,43 +36,28 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # Normalize the image (ImageNet mean and std)
 ])
 
-# Function to convert image from URL to PIL image
-def url_to_image(image_url):
-    response = requests.get(image_url)
-    img = Image.open(BytesIO(response.content))
-    return img
 
-# Function to convert base64 string to PIL image
-def base64_to_pil(base64_string):
-    img_data = base64.b64decode(base64_string)
-    return Image.open(BytesIO(img_data))
+def predict(image):
+    # Check if the input contains a base64-encoded string
+    if isinstance(image, dict) and image.get("data"):
+        # Decode the base64 string into a PIL image
+        image_data = base64.b64decode(image["data"])
+        image = Image.open(BytesIO(image_data))
 
-# Define the prediction function
-def predict(image_input):
-    # If input is a string (URL or base64 encoded), handle accordingly
-    if isinstance(image_input, str):
-        if image_input.startswith("http"):  # If URL
-            image = url_to_image(image_input)
-        elif image_input.startswith("data:image"):  # If base64 string
-            image = base64_to_pil(image_input)
-        else:  # Local image path
-            image = Image.open(image_input)
-    else:
-        image = image_input  # If the input is already a PIL image
+    # Apply your existing transformations
+    image = transform(image).unsqueeze(0)  # Transform and add batch dimension
+    image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
-    # Apply the necessary transformations to the image
-    image = transform(image).unsqueeze(0)  # Add batch dimension
-    image = image.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))  # Move to GPU if available
-    
+    # Perform inference
     with torch.no_grad():
-        outputs = model(image)  # Perform forward pass
-        predicted_class = torch.argmax(outputs, dim=1).item()  # Get the predicted class
+        outputs = model(image)
+        predicted_class = torch.argmax(outputs, dim=1).item()
     
     # 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:
-        return "The photo you've sent is of a healthy wheat image."
+        return "The photo you've sent is of a healthy maize image."
     else:
         return "Unexpected class prediction."