Debug: Issue with prediction

app.py

@@ -17,23 +17,34 @@ def load_model(model_path: str):
     """
     Load the model.
     """
-    # Check if CUDA is available and set device
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print(f"Using device: {device}")
     
-    # Load the pre-trained ResNet50 model
-    model = models.resnet50(weights=None)
+    # Load the model with default weights first
+    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
     model = model.to(device)
 
-    # Load custom weights from a .pth file
+    # Load custom weights
     state_dict = torch.load(model_path, map_location=device)
-
-    # Filter out unexpected keys
+    
+    # Debug: Print state dict info
+    print("\nState dict keys:", list(state_dict['model_state_dict'].keys())[:5])
+    print("Model state dict keys:", list(model.state_dict().keys())[:5])
+    
+    # Check if the final layer weights match
+    fc_weight_shape = state_dict['model_state_dict']['fc.weight'].shape
+    print(f"\nFC layer weight shape: {fc_weight_shape}")
+    
     filtered_state_dict = {k: v for k, v in state_dict['model_state_dict'].items() if k in model.state_dict()}
-
-    # Load the filtered state dictionary into the model
+    print(f"Filtered state dict size: {len(filtered_state_dict)} / {len(state_dict['model_state_dict'])}")
+    
     model.load_state_dict(filtered_state_dict, strict=False)
     model.eval()
+    
+    # Verify model
+    print("\nModel architecture:")
+    print(model)
+    
     return model
 
 
@@ -41,8 +52,8 @@ def load_classes():
     """
     Load the classes.
     """
-    # Get ImageNet class names from ResNet50 weights
-    weights = models.ResNet50_Weights.IMAGENET1K_V2
+    # Load classes from the same weights version as the model was trained with
+    weights = models.ResNet50_Weights.IMAGENET1K_V1  # Try V1 instead of V2
     classes = weights.meta["categories"]
     return classes
 
@@ -53,7 +64,7 @@ def inference_wrapper(image, alpha, top_k, target_layer):
     """
     try:
         if image is None:
-            return {"No image provided": 1.0}, None
+            return {"Error": 1.0}, None
             
         results = inference(
             image, 
@@ -65,19 +76,21 @@ def inference_wrapper(image, alpha, top_k, target_layer):
         )
         
         if results is None:
-            return {"Processing failed": 1.0}, None
+            return {"Error": 1.0}, None
             
         return results
         
-    except Exception as e:
+    except RuntimeError as e:
         error_msg = str(e)
         print(f"Error in inference: {error_msg}")
         
-        # Handle GPU quota error specifically
-        if "GPU quota" in error_msg:
-            return {"GPU quota exceeded - Please try again later": 1.0}, None
+        if "out of memory" in error_msg.lower():
+            return {"GPU Memory Error - Please try again": 1.0}, None
+        return {"Runtime Error: " + error_msg: 1.0}, None
         
-        # Handle other errors
+    except Exception as e:
+        error_msg = str(e)
+        print(f"Error in inference: {error_msg}")
         return {"Error: " + error_msg: 1.0}, None
 
 
@@ -166,12 +179,12 @@ def main():
 
                 # Examples section for Gradio 5.x
                 examples = [
-                    [
-                        "assets/examples/dog.jpg", 
-                        0.5,  # alpha slider
-                        3,    # top_k slider
-                        4     # target_layer slider
-                    ],
+                    # [
+                    #     "assets/examples/dog.jpg", 
+                    #     0.5,  # alpha slider
+                    #     3,    # top_k slider
+                    #     4     # target_layer slider
+                    # ],
                     [
                         "assets/examples/cat.jpg",
                         0.5,
@@ -190,12 +203,12 @@ def main():
                         3,
                         4
                     ],
-                    [
-                        "assets/examples/shark-plane.jpg",
-                        0.5,
-                        3,
-                        4
-                    ],
+                    # [
+                    #     "assets/examples/shark-plane.jpg",
+                    #     0.5,
+                    #     3,
+                    #     4
+                    # ],
                     [
                         "assets/examples/car.jpg",
                         0.5,
@@ -241,14 +254,14 @@ def main():
                         gradcam_output
                     ],
                     fn=inference_wrapper,
-                    cache_examples=True,
+                    cache_examples=False,  # Disable caching to prevent memory issues
                     label="Click on any example to run GradCAM"
                 )
 
             # Queue configuration
-            demo.queue(max_size=1)
+            demo.queue(max_size=1)  # Only allow one job at a time
 
-            # Launch configuration
+            # Launch with minimal memory usage
             demo.launch(
                 server_name="0.0.0.0",
                 server_port=7860,

inference.py

@@ -23,76 +23,87 @@ def inference(image, alpha, top_k, target_layer, model=None, classes=None):
     Run inference with GradCAM visualization
     """
     try:
+        # Clear CUDA cache before starting
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
         # Ensure model is on correct device and in eval mode
         model = model.to(device)
         model.eval()
         
-        # Save a copy of input img
-        org_img = image.copy()
+        # Process with reduced precision to save memory
+        with torch.cuda.amp.autocast():
+            # Save a copy of input img
+            org_img = image.copy()
 
-        # Convert img to tensor and normalize it
-        _transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(
-                mean=[0.485, 0.456, 0.406],
-                std=[0.229, 0.224, 0.225]
-            )
-        ])
+            # Convert img to tensor and normalize it
+            _transform = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406],
+                    std=[0.229, 0.224, 0.225]
+                )
+            ])
 
-        # Preprocess the input image and move to device
-        input_tensor = _transform(image).to(device)
-        input_tensor = input_tensor.unsqueeze(0)
-        input_tensor.requires_grad = True
-        
-        # Get Model Predictions
-        outputs = model(input_tensor)
-        probabilities = torch.softmax(outputs, dim=1)[0]
-        confidences = {classes[i]: float(probabilities[i]) for i in range(1000)}
+            # Preprocess the input image and move to device
+            input_tensor = _transform(image).to(device)
+            input_tensor = input_tensor.unsqueeze(0)
+            input_tensor.requires_grad = True
+            
+            # Get Model Predictions
+            outputs = model(input_tensor)
+            probabilities = torch.softmax(outputs, dim=1)[0]
+            confidences = {classes[i]: float(probabilities[i]) for i in range(1000)}
 
-        # Select the top classes based on user input
-        sorted_confidences = sorted(confidences.items(), key=lambda val: val[1], reverse=True)
-        show_confidences = OrderedDict(sorted_confidences[:top_k])
+            # Select the top classes based on user input
+            sorted_confidences = sorted(confidences.items(), key=lambda val: val[1], reverse=True)
+            show_confidences = OrderedDict(sorted_confidences[:top_k])
 
-        # Map layer numbers to meaningful parts of the ResNet architecture
-        _layers = {
-            1: model.conv1,
-            2: model.layer1[-1],
-            3: model.layer2[-1],
-            4: model.layer3[-1],
-            5: model.layer4[-1],
-            6: model.layer4[-1]
-        }
+            # Map layer numbers to meaningful parts of the ResNet architecture
+            _layers = {
+                1: model.conv1,
+                2: model.layer1[-1],
+                3: model.layer2[-1],
+                4: model.layer3[-1],
+                5: model.layer4[-1],
+                6: model.layer4[-1]
+            }
 
-        # Ensure valid layer selection
-        target_layer = min(max(target_layer, 1), 6)
-        target_layers = [_layers[target_layer]]
+            # Ensure valid layer selection
+            target_layer = min(max(target_layer, 1), 6)
+            target_layers = [_layers[target_layer]]
 
-        # Get the class activations from the selected layer
-        cam = GradCAM(model=model, target_layers=target_layers)
+            # Get the class activations from the selected layer
+            cam = GradCAM(model=model, target_layers=target_layers)
 
-        # Get the most probable class index
-        top_class = max(confidences.items(), key=lambda x: x[1])[0]
-        class_idx = classes.index(top_class)
-        
-        # Generate GradCAM for the top predicted class
-        grayscale_cam = cam(
-            input_tensor=input_tensor,
-            targets=[ClassifierOutputTarget(class_idx)],
-            aug_smooth=False,  # Disable augmentation for memory efficiency
-            eigen_smooth=False  # Disable eigen smoothing for memory efficiency
-        )
-        grayscale_cam = grayscale_cam[0, :]
+            # Get the most probable class index
+            top_class = max(confidences.items(), key=lambda x: x[1])[0]
+            class_idx = classes.index(top_class)
+            
+            # Generate GradCAM for the top predicted class
+            grayscale_cam = cam(
+                input_tensor=input_tensor,
+                targets=[ClassifierOutputTarget(class_idx)],
+                aug_smooth=False,  # Disable augmentation for memory efficiency
+                eigen_smooth=False  # Disable eigen smoothing for memory efficiency
+            )
+            grayscale_cam = grayscale_cam[0, :]
 
-        # Overlay input image with Class activations
-        visualization = show_cam_on_image(org_img/255., grayscale_cam, use_rgb=True, image_weight=alpha)
-        
-        # Clear CUDA cache
+            # Overlay input image with Class activations
+            visualization = show_cam_on_image(org_img/255., grayscale_cam, use_rgb=True, image_weight=alpha)
+            
+            # Clear CUDA cache after processing
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                
+            return show_confidences, visualization
+            
+    except torch.cuda.OutOfMemoryError:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-            
-        return show_confidences, visualization
+        raise RuntimeError("GPU out of memory - Please try again with a smaller image")
         
     except Exception as e:
         if torch.cuda.is_available():