Update app.py

app.py

@@ -12,6 +12,25 @@ from ultralytics import YOLO
 import numpy as np
 import cv2
 from modelscope.hub.snapshot_download import snapshot_download
+from ultralytics.nn.modules import Conv, C2f
+from torch import nn
+import ultralytics.nn.modules as modules
+
+# Add custom C3k2 module definition
+class C3k2(nn.Module):
+    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
+        super().__init__()
+        c_ = int(c2 * e)
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c1, c_, 1, 1)
+        self.cv3 = Conv(2 * c_, c2, 1)
+        self.m = nn.Sequential(*(C2f(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
+
+    def forward(self, x):
+        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
+
+# Patch the Ultralytics module
+modules.C3k2 = C3k2
 
 # Fix GLIBCXX dependency
 os.environ['LD_LIBRARY_PATH'] = '/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH'
@@ -34,7 +53,7 @@ MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
 try:
     model_dir = snapshot_download(MODEL_NAME, 
                                  cache_dir='./models',
-                                 revision='v1.0.0')  # Verified working revision
+                                 revision='main')
 except Exception as e:
     raise RuntimeError(f"Model download failed: {str(e)}")
 
@@ -56,7 +75,7 @@ def is_video(filename):
 
 @spaces.GPU
 def load_model_and_tokenizer():
-    """Load 8-bit quantized model with memory optimizations"""
+    """Load 8-bit quantized model"""
     try:
         torch.cuda.empty_cache()
         gc.collect()
@@ -82,7 +101,7 @@ def load_model_and_tokenizer():
         raise
 
 def process_yolo_results(results):
-    """Process YOLO detection results with safety checks"""
+    """Process YOLO detection results"""
     machinery_mapping = {
         'tower_crane': "Tower Crane",
         'mobile_crane': "Mobile Crane",
@@ -103,29 +122,26 @@ def process_yolo_results(results):
 
     counts = {"Worker": 0, **{v: 0 for v in machinery_mapping.values()}}
     
-    try:
-        for r in results:
-            for box in r.boxes:
-                if box.conf.item() < 0.5:
-                    continue
+    for r in results:
+        for box in r.boxes:
+            if box.conf.item() < 0.5:
+                continue
+            
+            cls_name = YOLO_MODEL.names[int(box.cls.item())].lower()
+            if cls_name == 'worker':
+                counts["Worker"] += 1
+                continue
                 
-                cls_name = YOLO_MODEL.names[int(box.cls.item())].lower()
-                if cls_name == 'worker':
-                    counts["Worker"] += 1
-                    continue
-                    
-                for key, value in machinery_mapping.items():
-                    if key in cls_name:
-                        counts[value] += 1
-                        break
-    except Exception as e:
-        print(f"YOLO processing error: {str(e)}")
-    
+            for key, value in machinery_mapping.items():
+                if key in cls_name:
+                    counts[value] += 1
+                    break
+
     return counts["Worker"], sum(counts.values()) - counts["Worker"], counts
 
 @spaces.GPU
 def detect_people_and_machinery(media_path):
-    """GPU-accelerated detection with memory management"""
+    """GPU-accelerated detection"""
     try:
         max_people = 0
         max_machines = {k: 0 for k in [
@@ -169,7 +185,7 @@ def detect_people_and_machinery(media_path):
 
 @spaces.GPU
 def analyze_video_activities(video_path):
-    """Video analysis with chunk processing and memory cleanup"""
+    """Video analysis with chunk processing"""
     try:
         model, tokenizer, processor = load_model_and_tokenizer()
         responses = []
@@ -178,7 +194,6 @@ def analyze_video_activities(video_path):
         frame_step = max(1, int(vr.get_avg_fps()))
         total_frames = len(vr)
         
-        # Process in 16-frame chunks
         for i in range(0, total_frames, 16):
             end_idx = min(i+16, total_frames)
             frames = [Image.fromarray(vr[j].asnumpy()) for j in range(i, end_idx)]
@@ -203,7 +218,7 @@ def analyze_video_activities(video_path):
 
 @spaces.GPU
 def analyze_image_activities(image_path):
-    """Image analysis with memory cleanup"""
+    """Image analysis pipeline"""
     try:
         model, tokenizer, processor = load_model_and_tokenizer()
         image = Image.open(image_path).convert("RGB")
@@ -225,7 +240,7 @@ def analyze_image_activities(image_path):
 
 @spaces.GPU
 def annotate_video_with_bboxes(video_path):
-    """Video annotation with efficient frame processing"""
+    """Video annotation with detection overlay"""
     try:
         cap = cv2.VideoCapture(video_path)
         fps = cap.get(cv2.CAP_PROP_FPS)
@@ -241,7 +256,6 @@ def annotate_video_with_bboxes(video_path):
             if not ret:
                 break
             
-            # Process every 5th frame to reduce load
             if frame_count % 5 == 0:
                 results = YOLO_MODEL(frame)
                 counts = {}
@@ -275,7 +289,7 @@ def annotate_video_with_bboxes(video_path):
         return None
 
 def process_diary(day, date, media):
-    """Main processing pipeline with error handling"""
+    """Main processing pipeline"""
     try:
         if not media:
             return [day, date, "No data", "No data", "No data", "No data", None]