Fix: Prevent Silent Exception Swallowing in Model Loading

Author: priyan17singh

perceptionmetrics/models/torch_detection.py

@@ -263,23 +263,22 @@ def __init__(
 
         # Load model from file or use passed instance
         if isinstance(model, str):
-            assert os.path.isfile(model), "Torch model file not found"
+            if not os.path.isfile(model):
+                raise FileNotFoundError(f"Model file not found: {model}")
             model_fname = model
             try:
                 model = torch.jit.load(model, map_location=self.device)
                 model_type = "compiled"
-            except RuntimeError:
+            except (RuntimeError, EOFError) as jit_err:
+                # TorchScript load failed, try loading as native PyTorch
                 try:
-                    model = torch.load(
-                        model, map_location=self.device, weights_only=False
-                    )
+                    model = torch.load(model, map_location=self.device, weights_only=False)
                     model_type = "native"
-                except Exception as e:
-                    raise ValueError(
-                        f"Failed to load model. "
-                        f"Ensure it is a valid PyTorch or TorchScript model. Error : {e}"
-                    )
-
+                except Exception as load_err:
+                    raise RuntimeError(
+                        f"Failed to load model as TorchScript or PyTorch module. "
+                        f"TorchScript error: {jit_err}. PyTorch error: {load_err}"
+                    ) from load_err
         elif isinstance(model, torch.nn.Module):
             model_fname = None
             model_type = "native"

perceptionmetrics/models/torch_segmentation.py

@@ -217,15 +217,21 @@ def __init__(
 
         # If 'model' contains a string, check that it is a valid filename and load model
         if isinstance(model, str):
-            assert os.path.isfile(model), "TorchScript Model file not found"
+            if not os.path.isfile(model):
+                raise FileNotFoundError(f"Model file not found: {model}")
             model_fname = model
             try:
                 model = torch.jit.load(model, map_location=self.device)
                 model_type = "compiled"
-            except:
-                print("Model is not a TorchScript model. Loading as a PyTorch module.")
-                model = torch.load(model, map_location=self.device)
-                model_type = "native"
+            except (RuntimeError, EOFError) as jit_err:
+                try:
+                    model = torch.load(model, map_location=self.device, weights_only=False)
+                    model_type = "native"
+                except Exception as load_err:
+                    raise RuntimeError(
+                        f"Failed to load model as TorchScript or PyTorch module. "
+                        f"TorchScript error: {jit_err}. PyTorch error: {load_err}"
+                    ) from load_err
         # Otherwise, check that it is a PyTorch module
         elif isinstance(model, torch.nn.Module):
             model_fname = None
@@ -563,15 +569,22 @@ def __init__(
 
         # If 'model' contains a string, check that it is a valid filename and load model
         if isinstance(model, str):
-            assert os.path.isfile(model), "TorchScript Model file not found"
+            if not os.path.isfile(model):
+                raise FileNotFoundError(f"Model file not found: {model}")
             model_fname = model
             try:
                 model = torch.jit.load(model, map_location=self.device)
                 model_type = "compiled"
-            except Exception:
-                print("Model is not a TorchScript model. Loading as a PyTorch module.")
-                model = torch.load(model, map_location=self.device)
-                model_type = "native"
+            except (RuntimeError, EOFError) as jit_err:
+                # TorchScript load failed, try loading as native PyTorch
+                try:
+                    model = torch.load(model, map_location=self.device, weights_only=False)
+                    model_type = "native"
+                except Exception as load_err:
+                    raise RuntimeError(
+                        f"Failed to load model as TorchScript or PyTorch module. "
+                        f"TorchScript error: {jit_err}. PyTorch error: {load_err}"
+                    ) from load_err
 
         # Otherwise, check that it is a PyTorch module
         elif isinstance(model, torch.nn.Module):

perceptionmetrics/models/utils/o3d/kpconv.py

@@ -4,8 +4,8 @@
 
 try:
     from open3d._ml3d.torch.models.kpconv import batch_grid_subsampling, batch_neighbors
-except Exception:
-    print("Open3D-ML3D not available")
+except ImportError:
+    print("Warning: Open3D-ML3D not available. KPConv functionality will be limited.")
 import torch
 
 import perceptionmetrics.utils.lidar as ul

perceptionmetrics/models/utils/o3d/randlanet.py

@@ -4,8 +4,8 @@
 
 try:
     from open3d._ml3d.datasets.utils import DataProcessing
-except Exception:
-    print("Open3D-ML3D not available")
+except ImportError:
+    print("Warning: Open3D-ML3D not available. RandLANet functionality will be limited.")
 import torch
 
 import perceptionmetrics.utils.lidar as ul

tests/test_model_loading.py

@@ -0,0 +1,99 @@
+"""Unit tests for model loading robustness"""
+import os
+import tempfile
+import pytest
+import torch
+from PIL import Image
+
+from perceptionmetrics.models.torch_detection import TorchImageDetectionModel
+from perceptionmetrics.models.torch_segmentation import TorchImageSegmentationModel
+
+
+class TestModelLoadingExceptions:
+    """Test that model loading raises appropriate exceptions"""
+    
+    @pytest.fixture
+    def temp_files(self):
+        """Create temporary test files"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create dummy model file (corrupted)
+            bad_model = os.path.join(tmpdir, "bad_model.pt")
+            with open(bad_model, "w") as f:
+                f.write("this is not a pytorch model")
+            
+            # Create dummy ontology
+            ontology = os.path.join(tmpdir, "ontology.json")
+            import json
+            with open(ontology, "w") as f:
+                json.dump({
+                    "car": {"idx": 0, "rgb": [0, 0, 0]},
+                    "person": {"idx": 1, "rgb": [255, 0, 0]},
+                }, f)
+            
+            # Create dummy config
+            config = os.path.join(tmpdir, "config.json")
+            with open(config, "w") as f:
+                json.dump({
+                    "resize": {"width": 512, "height": 512},
+                    "normalization": {
+                        "mean": [0.485, 0.456, 0.406],
+                        "std": [0.229, 0.224, 0.225]
+                    },
+                    "batch_size": 1,
+                    "model_format": "torchvision"
+                }, f)
+            
+            yield {
+                "bad_model": bad_model,
+                "ontology": ontology,
+                "config": config,
+                "tmpdir": tmpdir
+            }
+    
+    def test_detection_model_bad_file_raises_specific_error(self, temp_files):
+        """Test that loading corrupted model raises RuntimeError, not generic Exception"""
+        with pytest.raises(RuntimeError) as exc_info:
+            TorchImageDetectionModel(
+                model=temp_files["bad_model"],
+                model_cfg=temp_files["config"],
+                ontology_fname=temp_files["ontology"]
+            )
+        
+        # Check error message is informative
+        error_msg = str(exc_info.value)
+        assert "Failed to load model" in error_msg
+        assert "TorchScript error" in error_msg or "PyTorch error" in error_msg
+    
+    def test_segmentation_model_bad_file_raises_specific_error(self, temp_files):
+        """Test that loading corrupted segmentation model raises RuntimeError"""
+        with pytest.raises(RuntimeError) as exc_info:
+            TorchImageSegmentationModel(
+                model=temp_files["bad_model"],
+                model_cfg=temp_files["config"],
+                ontology_fname=temp_files["ontology"]
+            )
+        
+        error_msg = str(exc_info.value)
+        assert "Failed to load model" in error_msg
+    
+    def test_detection_model_missing_file_raises_file_not_found(self, temp_files):
+        """Test that missing model file raises FileNotFoundError"""
+        with pytest.raises(FileNotFoundError) as exc_info:
+            TorchImageDetectionModel(
+                model="/nonexistent/path/model.pt",
+                model_cfg=temp_files["config"],
+                ontology_fname=temp_files["ontology"]
+            )
+        
+        assert "Model file not found" in str(exc_info.value)
+    
+    def test_segmentation_model_missing_file_raises_file_not_found(self, temp_files):
+        """Test that missing segmentation model file raises FileNotFoundError"""
+        with pytest.raises(FileNotFoundError) as exc_info:
+            TorchImageSegmentationModel(
+                model="/nonexistent/path/model.pt",
+                model_cfg=temp_files["config"],
+                ontology_fname=temp_files["ontology"]
+            )
+        
+        assert "Model file not found" in str(exc_info.value)