Update app.py

app.py

@@ -34,15 +34,27 @@ class BehaviorDetector(nn.Module):
         return torch.sigmoid(self.base_model(x))
 
 class DogBehaviorAnalyzer:
-    def __init__(self, model_path='best.pt'):
+    def __init__(self, model_path=None):
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
-        # Initialize YOLO model for dog detection
-        self.yolo_model = YOLO(model_path) if model_path else None
+        # Initialize YOLO model for dog detection (optional)
+        try:
+            self.yolo_model = YOLO(model_path) if model_path else None
+        except Exception as e:
+            st.warning("YOLO model not found. Running without dog detection.")
+            self.yolo_model = None
         
         # Initialize behavior classifier
         self.num_behaviors = 5
-        self.behavior_model = BehaviorDetector(self.num_behaviors).to(self.device)
+        try:
+            self.behavior_model = BehaviorDetector(self.num_behaviors).to(self.device)
+        except Exception as e:
+            st.warning("Error loading behavior model. Using default classifier.")
+            self.behavior_model = models.resnet18(pretrained=True)
+            num_features = self.behavior_model.fc.in_features
+            self.behavior_model.fc = nn.Linear(num_features, self.num_behaviors)
+            self.behavior_model = self.behavior_model.to(self.device)
+        
         self.behavior_model.eval()
         
         # Define sophisticated transforms