Improve model loading with better error handling and fallback mechanisms

app.py

@@ -77,7 +77,6 @@ def load_models():
         }
         
         # Check if model files exist
-        import os
         if not os.path.exists('mobilenet0.25_Final.pth'):
             print("❌ mobilenet0.25_Final.pth not found!")
             return False
@@ -87,28 +86,69 @@ def load_models():
             
         print("Model files found, loading MobileNet...")
         
-        # Load MobileNet model
-        mobilenet_model = RetinaFace(cfg=mobilenet_cfg, phase='test')
-        mobilenet_model.load_state_dict(torch.load('mobilenet0.25_Final.pth', map_location=device))
-        mobilenet_model.eval()
-        mobilenet_model = mobilenet_model.to(device)
-        print("✅ MobileNet model loaded successfully!")
+        # Load MobileNet model with better error handling
+        try:
+            mobilenet_model = RetinaFace(cfg=mobilenet_cfg, phase='test')
+            print("✅ MobileNet model instance created")
+            
+            # Load state dict
+            mobilenet_state = torch.load('mobilenet0.25_Final.pth', map_location=device)
+            print(f"✅ MobileNet state dict loaded with {len(mobilenet_state.keys())} keys")
+            
+            # Try to load state dict with strict=False to handle key mismatches
+            missing_keys, unexpected_keys = mobilenet_model.load_state_dict(mobilenet_state, strict=False)
+            
+            if missing_keys:
+                print(f"⚠️ Missing keys in MobileNet: {missing_keys[:5]}...")  # Show first 5
+            if unexpected_keys:
+                print(f"⚠️ Unexpected keys in MobileNet: {unexpected_keys[:5]}...")  # Show first 5
+                
+            mobilenet_model.eval()
+            mobilenet_model = mobilenet_model.to(device)
+            print("✅ MobileNet model loaded successfully!")
+            
+        except Exception as e:
+            print(f"❌ Error loading MobileNet: {e}")
+            mobilenet_model = None
         
         print("Loading ResNet...")
         
-        # Load ResNet model
-        resnet_model = RetinaFace(cfg=resnet_cfg, phase='test')
-        resnet_model.load_state_dict(torch.load('Resnet50_Final.pth', map_location=device))
-        resnet_model.eval()
-        resnet_model = resnet_model.to(device)
-        print("✅ ResNet model loaded successfully!")
+        # Load ResNet model with better error handling
+        try:
+            resnet_model = RetinaFace(cfg=resnet_cfg, phase='test')
+            print("✅ ResNet model instance created")
+            
+            # Load state dict
+            resnet_state = torch.load('Resnet50_Final.pth', map_location=device)
+            print(f"✅ ResNet state dict loaded with {len(resnet_state.keys())} keys")
+            
+            # Try to load state dict with strict=False to handle key mismatches
+            missing_keys, unexpected_keys = resnet_model.load_state_dict(resnet_state, strict=False)
+            
+            if missing_keys:
+                print(f"⚠️ Missing keys in ResNet: {missing_keys[:5]}...")  # Show first 5
+            if unexpected_keys:
+                print(f"⚠️ Unexpected keys in ResNet: {unexpected_keys[:5]}...")  # Show first 5
+                
+            resnet_model.eval()
+            resnet_model = resnet_model.to(device)
+            print("✅ ResNet model loaded successfully!")
+            
+        except Exception as e:
+            print(f"❌ Error loading ResNet: {e}")
+            resnet_model = None
         
-        print("✅ All models loaded successfully!")
-        return True
+        # Check if at least one model loaded
+        if mobilenet_model is not None or resnet_model is not None:
+            print("✅ At least one model loaded successfully!")
+            return True
+        else:
+            print("❌ No models loaded successfully!")
+            return False
         
     except Exception as e:
         import traceback
-        print(f"❌ Error loading models: {e}")
+        print(f"❌ Error in load_models: {e}")
         print(f"❌ Full traceback: {traceback.format_exc()}")
         return False
 
@@ -127,6 +167,12 @@ def detect_faces(image, model_type="mobilenet", confidence_threshold=0.5, nms_th
                 'clip': False,
                 'image_size': 840
             }
+            if model is None:
+                # Fallback to MobileNet if ResNet not available
+                print("⚠️ ResNet not available, falling back to MobileNet")
+                model = mobilenet_model
+                model_type = "mobilenet"
+                cfg['image_size'] = 640
         else:
             model = mobilenet_model
             cfg = {
@@ -136,9 +182,15 @@ def detect_faces(image, model_type="mobilenet", confidence_threshold=0.5, nms_th
                 'clip': False,
                 'image_size': 640
             }
+            if model is None:
+                # Fallback to ResNet if MobileNet not available
+                print("⚠️ MobileNet not available, falling back to ResNet")
+                model = resnet_model
+                model_type = "resnet"
+                cfg['image_size'] = 840
         
         if model is None:
-            return None, "Models not loaded"
+            return None, "❌ No models are loaded. Please check the model loading logs."
             
         # Convert PIL to numpy array
         if isinstance(image, Image.Image):