Fix: Runtime error

app.py

@@ -67,7 +67,7 @@ def inference_wrapper(image, alpha, top_k, target_layer):
                 )
     except Exception as e:
         print(f"Error in inference: {str(e)}")
-        return gr.Error(f"Error processing image: {str(e)}")
+        return None, None
 
 
 def main():

inference.py

@@ -28,35 +28,23 @@ def inference(image, alpha, top_k, target_layer, model=None, classes=None):
     model = model.to(device)
     model.eval()
     
-    # Convert input to tensor and move to GPU
-    if isinstance(image, np.ndarray):
-        image_tensor = torch.from_numpy(image).to(device)
-        if image_tensor.ndim == 3:
-            image_tensor = image_tensor.unsqueeze(0)
-    else:
-        image_tensor = image.to(device)
-    
     with torch.cuda.amp.autocast():  # Enable automatic mixed precision
         with torch.no_grad():
             # Save a copy of input img
             org_img = image.copy()
 
-            # Calculate mean over each channel of input image
+            # Calculate mean and std over each channel
             mean_r, mean_g, mean_b = np.mean(image[:, :, 0]/255.), np.mean(image[:, :, 1]/255.), np.mean(image[:, :, 2]/255.)
-
-            # Calculate Standard deviation over each channel
             std_r, std_g, std_b = np.std(image[:, :, 0]/255.), np.std(image[:, :, 1]/255.), np.std(image[:, :, 2]/255.)
 
             # Convert img to tensor and normalize it
             _transform = transforms.Compose([
                 transforms.ToTensor(),
                 transforms.Normalize((mean_r, mean_g, mean_b), (std_r, std_g, std_b))
-                ])
-
-            # Preprocess the input image
-            input_tensor = _transform(image)
+            ])
 
-            # Create a mini-batch as expected by the model
+            # Preprocess the input image and move to device
+            input_tensor = _transform(image).to(device)
             input_tensor = input_tensor.unsqueeze(0)
 
             # Get Model Predictions