Update app.py

app.py

@@ -6,45 +6,35 @@ import time
 from ultralytics import YOLO
 import spaces
 import os
-import logging
-@spaces.GPU
 
 class CrowdDetection:
     def __init__(self, model_path="yolov8n.pt"):
-        logger.info(f"Initializing CrowdDetection with model: {model_path}")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        try:
-            if not os.path.exists(model_path):
-                logger.info(f"Model {model_path} not found, downloading...")
-                self.model = YOLO("yolov8n.pt")  # Downloads if not present
-                self.model.save(model_path)
-            else:
-                self.model = YOLO(model_path)
-            self.model.to(self.device)
-            logger.info("CrowdDetection model loaded successfully")
-        except Exception as e:
-            logger.error(f"Failed to initialize model: {str(e)}")
-            raise
+        self.model_path = model_path  # Store path, load model later
 
+    @spaces.GPU
     def detect_crowd(self, video_path):
-        logger.info(f"Processing video for crowd detection: {video_path}")
         try:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if not os.path.exists(self.model_path):
+                model = YOLO("yolov8n.pt")
+                model.save(self.model_path)
+            else:
+                model = YOLO(self.model_path)
+            model.to(device)
+
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
-                logger.error(f"Failed to open video: {video_path}")
                 raise ValueError(f"❌ Failed to open video: {video_path}")
 
             fps = int(cap.get(cv2.CAP_PROP_FPS))
             width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
             height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            logger.debug(f"Video specs - FPS: {fps}, Width: {width}, Height: {height}")
 
             output_path = "output_crowd.mp4"
             fourcc = cv2.VideoWriter_fourcc(*"mp4v")
             out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
             if not out.isOpened():
                 cap.release()
-                logger.error(f"Failed to initialize video writer for {output_path}")
                 raise ValueError(f"❌ Failed to initialize video writer")
 
             CROWD_THRESHOLD = 10
@@ -56,9 +46,8 @@ class CrowdDetection:
                     break
                 frame_count += 1
 
-                results = self.model(frame)
+                results = model(frame)
                 person_count = sum(1 for result in results for cls in result.boxes.cls.cpu().numpy() if int(cls) == 0)
-                logger.debug(f"Frame {frame_count}: Detected {person_count} people")
 
                 for result in results:
                     boxes = result.boxes.xyxy.cpu().numpy()
@@ -77,28 +66,26 @@ class CrowdDetection:
             cap.release()
             out.release()
             if frame_count == 0 or not os.path.exists(output_path):
-                logger.error(f"Processing failed: Frames processed: {frame_count}, Output exists: {os.path.exists(output_path)}")
                 raise ValueError("❌ Processing failed: No frames processed or output not created")
-            logger.info(f"Crowd detection completed, output saved to: {output_path}")
             return output_path
         except Exception as e:
-            logger.error(f"Error in detect_crowd: {str(e)}")
-            raise
+            raise ValueError(f"Error in detect_crowd: {str(e)}")
 
 class PeopleTracking:
     def __init__(self, yolo_model_path="yolov8n.pt"):
-        logger.info(f"Initializing PeopleTracking with model: {yolo_model_path}")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if not os.path.exists(yolo_model_path):
-            self.model = YOLO("yolov8n.pt")
-            self.model.save(yolo_model_path)
-        else:
-            self.model = YOLO(yolo_model_path)
-        self.model.to(self.device)
+        self.model_path = yolo_model_path
 
+    @spaces.GPU
     def track_people(self, video_path):
-        logger.info(f"Tracking people in video: {video_path}")
         try:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if not os.path.exists(self.model_path):
+                model = YOLO("yolov8n.pt")
+                model.save(self.model_path)
+            else:
+                model = YOLO(self.model_path)
+            model.to(device)
+
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
                 raise ValueError(f"❌ Failed to open video: {video_path}")
@@ -117,7 +104,7 @@ class PeopleTracking:
                 if not ret:
                     break
 
-                results = self.model.track(frame, persist=True)
+                results = model.track(frame, persist=True)
                 for result in results:
                     boxes = result.boxes.xyxy.cpu().numpy()
                     classes = result.boxes.cls.cpu().numpy()
@@ -137,23 +124,23 @@ class PeopleTracking:
                 raise ValueError("❌ Processing failed")
             return output_path
         except Exception as e:
-            logger.error(f"Error in track_people: {str(e)}")
-            raise
+            raise ValueError(f"Error in track_people: {str(e)}")
 
 class FallDetection:
     def __init__(self, yolo_model_path="yolov8l.pt"):
-        logger.info(f"Initializing FallDetection with model: {yolo_model_path}")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if not os.path.exists(yolo_model_path):
-            self.model = YOLO("yolov8l.pt")
-            self.model.save(yolo_model_path)
-        else:
-            self.model = YOLO(yolo_model_path)
-        self.model.to(self.device)
+        self.model_path = yolo_model_path
 
+    @spaces.GPU
     def detect_fall(self, video_path):
-        logger.info(f"Detecting falls in video: {video_path}")
         try:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if not os.path.exists(self.model_path):
+                model = YOLO("yolov8l.pt")
+                model.save(self.model_path)
+            else:
+                model = YOLO(self.model_path)
+            model.to(device)
+
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
                 raise ValueError(f"❌ Failed to open video: {video_path}")
@@ -172,7 +159,7 @@ class FallDetection:
                 if not ret:
                     break
 
-                results = self.model(frame)
+                results = model(frame)
                 for result in results:
                     boxes = result.boxes.xyxy.cpu().numpy()
                     classes = result.boxes.cls.cpu().numpy()
@@ -202,23 +189,23 @@ class FallDetection:
                 raise ValueError("❌ Processing failed")
             return output_path
         except Exception as e:
-            logger.error(f"Error in detect_fall: {str(e)}")
-            raise
+            raise ValueError(f"Error in detect_fall: {str(e)}")
 
 class FightDetection:
     def __init__(self, yolo_model_path="yolov8n-pose.pt"):
-        logger.info(f"Initializing FightDetection with model: {yolo_model_path}")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if not os.path.exists(yolo_model_path):
-            self.model = YOLO("yolov8n-pose.pt")
-            self.model.save(yolo_model_path)
-        else:
-            self.model = YOLO(yolo_model_path)
-        self.model.to(self.device)
+        self.model_path = yolo_model_path
 
+    @spaces.GPU
     def detect_fight(self, video_path):
-        logger.info(f"Detecting fights in video: {video_path}")
         try:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if not os.path.exists(self.model_path):
+                model = YOLO("yolov8n-pose.pt")
+                model.save(self.model_path)
+            else:
+                model = YOLO(self.model_path)
+            model.to(device)
+
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
                 raise ValueError(f"❌ Failed to open video: {video_path}")
@@ -237,7 +224,7 @@ class FightDetection:
                 if not ret:
                     break
 
-                results = self.model.track(frame, persist=True)
+                results = model.track(frame, persist=True)
                 fight_detected = False
                 person_count = 0
 
@@ -267,8 +254,7 @@ class FightDetection:
                 raise ValueError("❌ Processing failed")
             return output_path
         except Exception as e:
-            logger.error(f"Error in detect_fight: {str(e)}")
-            raise
+            raise ValueError(f"Error in detect_fight: {str(e)}")
 
 # Unified processing function with status output
 def process_video(feature, video):
@@ -284,7 +270,6 @@ def process_video(feature, video):
         output_path = getattr(detector, method_name)(video)
         return f"{feature} completed successfully", output_path
     except Exception as e:
-        logger.error(f"Error processing video with {feature}: {str(e)}")
         return f"Error: {str(e)}", None
 
 # Gradio Interface with dual outputs