modifed app

app.py

@@ -72,49 +72,75 @@ class ModelPredictor:
         """
         Make prediction for a single image
         Args:
-            image: PIL Image or path to image
+            image: numpy array from Gradio
         Returns:
             Dictionary of class labels and probabilities
         """
-        if isinstance(image, str):
-            image = Image.open(image).convert("RGB")
-        else:
-            image = Image.fromarray(image).convert("RGB")
+        try:
+            # Convert numpy array to PIL Image
+            if isinstance(image, np.ndarray):
+                # If image is from Gradio, it will be a numpy array
+                image = Image.fromarray(image.astype("uint8"))
+            elif isinstance(image, str):
+                # If image is a file path
+                image = Image.open(image)
 
-        image_tensor = self.transform(image).unsqueeze(0)
-        image_tensor = image_tensor.to(self.device)
+            # Ensure image is in RGB mode
+            if image.mode != "RGB":
+                image = image.convert("RGB")
 
-        with torch.no_grad():
-            outputs = self.model(image_tensor)
-            probabilities = torch.nn.functional.softmax(outputs, dim=1)
+            # Apply transforms and predict
+            image_tensor = self.transform(image).unsqueeze(0)
+            image_tensor = image_tensor.to(self.device)
 
-            # Get top 5 predictions
-            top_probs, top_indices = torch.topk(probabilities, 5)
+            with torch.no_grad():
+                outputs = self.model(image_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)
 
-            # Create results dictionary
-            results = {}
-            for prob, idx in zip(top_probs[0], top_indices[0]):
-                class_name = self.class_labels[str(idx.item())]
-                results[class_name] = float(prob)
+                # Get top 5 predictions
+                top_probs, top_indices = torch.topk(probabilities, 5)
 
-        return results
+                # Create results dictionary
+                results = {}
+                for prob, idx in zip(top_probs[0], top_indices[0]):
+                    class_name = self.class_labels[str(idx.item())]
+                    results[class_name] = float(prob)
+
+            return results
+
+        except Exception as e:
+            print(f"Error in prediction: {str(e)}")
+            return {"error": 1.0}
 
 
 # Initialize the predictor
-predictor = ModelPredictor(
-    model_repo="Adityak204/ResNetVision-1K",  # Replace with your repo
-    model_filename="resnet50-epoch36-acc60.3506.ckpt",  # Replace with your model filename
-)
+try:
+    predictor = ModelPredictor(
+        model_repo="Adityak204/ResNetVision-1K",  # Replace with your repo
+        model_filename="resnet50-epoch36-acc60.3506.ckpt",  # Replace with your model filename
+    )
+except Exception as e:
+    print(f"Error initializing predictor: {str(e)}")
 
 
 def predict_image(image):
     """
     Gradio interface function
+    Args:
+        image: numpy array from Gradio's image input
+    Returns:
+        Dictionary of predictions formatted for display
     """
-    predictions = predictor.predict(image)
-
-    # Format results for display
-    return {k: f"{v:.2%}" for k, v in predictions.items()}
+    if image is None:
+        return {"Error: No image provided": 1.0}
+
+    try:
+        predictions = predictor.predict(image)
+        # Format results for display
+        return {k: f"{v:.2%}" for k, v in predictions.items()}
+    except Exception as e:
+        print(f"Error in predict_image: {str(e)}")
+        return {"Error: Failed to process image": 1.0}
 
 
 # Create Gradio interface