Update app.py

app.py

@@ -20,12 +20,14 @@ def initialize_gpu():
         torch.randn(10).cuda()
 initialize_gpu()
 
-# Load YOLO model
-YOLO_MODEL = YOLO('best_yolov11.pt')  # Keep this file in repo root
+# Load YOLO model with relative path
+YOLO_MODEL = YOLO('best_yolov11.pt')
 
-# Model configuration
+# Model configuration with quantization
 MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
-model_dir = snapshot_download(MODEL_NAME, cache_dir='./models')
+model_dir = snapshot_download(MODEL_NAME, 
+                             revision='v1.0.0',  # Specific revision
+                             cache_dir='./models')
 
 # Device setup
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -45,7 +47,7 @@ def is_video(filename):
 
 @spaces.GPU
 def load_model_and_tokenizer():
-    """Load 4-bit quantized model for memory efficiency"""
+    """Load 4-bit quantized model"""
     try:
         torch.cuda.empty_cache()
         gc.collect()
@@ -70,184 +72,10 @@ def load_model_and_tokenizer():
         print(f"Model loading error: {str(e)}")
         raise
 
-def process_yolo_results(results):
-    """Process YOLO detection results"""
-    machinery_mapping = {
-        'tower_crane': "Tower Crane",
-        'mobile_crane': "Mobile Crane",
-        'compactor': "Compactor/Roller",
-        'roller': "Compactor/Roller",
-        'bulldozer': "Bulldozer",
-        'dozer': "Bulldozer",
-        'excavator': "Excavator",
-        'dump_truck': "Dump Truck",
-        'truck': "Dump Truck",
-        'concrete_mixer_truck': "Concrete Mixer",
-        'loader': "Loader",
-        'pump_truck': "Pump Truck",
-        'pile_driver': "Pile Driver",
-        'grader': "Grader",
-        'other_vehicle': "Other Vehicle"
-    }
+# ... [Keep the rest of your existing functions unchanged] ...
 
-    counts = {"Worker": 0, **{v: 0 for v in machinery_mapping.values()}}
-    
-    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
-                
-            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"""
-    try:
-        max_people = 0
-        max_machines = {k: 0 for k in [
-            "Tower Crane", "Mobile Crane", "Compactor/Roller", "Bulldozer",
-            "Excavator", "Dump Truck", "Concrete Mixer", "Loader",
-            "Pump Truck", "Pile Driver", "Grader", "Other Vehicle"
-        ]}
-
-        if isinstance(media_path, str) and is_video(media_path):
-            cap = cv2.VideoCapture(media_path)
-            fps = cap.get(cv2.CAP_PROP_FPS)
-            sample_rate = max(1, int(fps))
-            
-            while cap.isOpened():
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                
-                results = YOLO_MODEL(frame)
-                people, machines, types = process_yolo_results(results)
-                
-                max_people = max(max_people, people)
-                for k in max_machines:
-                    max_machines[k] = max(max_machines[k], types.get(k, 0))
-            
-            cap.release()
-        else:
-            img = cv2.imread(media_path) if isinstance(media_path, str) else cv2.cvtColor(np.array(media_path), cv2.COLOR_RGB2BGR)
-            results = YOLO_MODEL(img)
-            max_people, _, types = process_yolo_results(results)
-            for k in max_machines:
-                max_machines[k] = types.get(k, 0)
-
-        filtered = {k: v for k, v in max_machines.items() if v > 0}
-        return max_people, sum(filtered.values()), filtered
-    
-    except Exception as e:
-        print(f"Detection error: {str(e)}")
-        return 0, 0, {}
-
-@spaces.GPU
-def analyze_video_activities(video_path):
-    """Video analysis with chunk processing"""
-    try:
-        model, tokenizer, processor = load_model_and_tokenizer()
-        responses = []
-        
-        vr = VideoReader(video_path, ctx=cpu(0))
-        frame_step = max(1, int(vr.get_avg_fps()))
-        frames = [Image.fromarray(f.asnumpy()) for f in vr[::frame_step]]
-        
-        # Process in chunks
-        for i in range(0, len(frames), 16):
-            chunk = frames[i:i+16]
-            inputs = processor(
-                [{"role": "user", "content": "Analyze construction activities", "video_frames": chunk}],
-                videos=[chunk]
-            ).to(DEVICE)
-            
-            response = model.generate(**inputs, max_new_tokens=200)
-            responses.append(response[0])
-            
-        del model, tokenizer, processor
-        torch.cuda.empty_cache()
-        return "\n".join(responses)
-    
-    except Exception as e:
-        print(f"Video analysis error: {str(e)}")
-        return "Activity analysis unavailable"
-
-@spaces.GPU
-def analyze_image_activities(image_path):
-    """Image analysis pipeline"""
-    try:
-        model, tokenizer, processor = load_model_and_tokenizer()
-        image = Image.open(image_path).convert("RGB")
-        
-        inputs = processor(
-            [{"role": "user", "content": "Analyze construction site", "images": [image]}],
-            images=[image]
-        ).to(DEVICE)
-        
-        response = model.generate(**inputs, max_new_tokens=200)
-        del model, tokenizer, processor
-        return response[0]
-    
-    except Exception as e:
-        print(f"Image analysis error: {str(e)}")
-        return "Activity analysis unavailable"
-
-@spaces.GPU
-def annotate_video_with_bboxes(video_path):
-    """Video annotation with real-time detection"""
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    writer = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        results = YOLO_MODEL(frame)
-        counts = {}
-        
-        for r in results:
-            for box in r.boxes:
-                if box.conf.item() < 0.5:
-                    continue
-                
-                cls_id = int(box.cls.item())
-                class_name = YOLO_MODEL.names[cls_id]
-                counts[class_name] = counts.get(class_name, 0) + 1
-                
-                # Draw bounding box
-                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
-                cv2.putText(frame, f"{class_name} {box.conf.item():.2f}", 
-                           (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
-        
-        # Add summary text
-        summary = ", ".join([f"{k}:{v}" for k,v in counts.items()])
-        cv2.putText(frame, summary, (10,30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,255), 2)
-        
-        writer.write(frame)
-    
-    cap.release()
-    writer.release()
-    return temp_file.name
-
-def process_diary(day, date, people, machinery, machinery_types, activities, media):
-    """Main processing pipeline"""
+def process_diary(day, date, media):
+    """Simplified processing pipeline"""
     try:
         if not media:
             return [day, date, "No data", "No data", "No data", "No data", None]
@@ -302,11 +130,19 @@ with gr.Blocks(title="Digital Site Diary", css="video {height: auto !important;}
             model_activities = gr.Textbox(label="Activity Analysis", lines=4)
             model_video = gr.Video(label="Safety Annotations")
 
+    # Fixed input mapping
     submit_btn.click(
         process_diary,
-        inputs=[day, date, None, None, None, None, media],
-        outputs=[model_day, model_date, model_people, model_machinery, 
-                model_machinery_types, model_activities, model_video]
+        inputs=[day, date, media],
+        outputs=[
+            model_day, 
+            model_date, 
+            model_people, 
+            model_machinery,
+            model_machinery_types, 
+            model_activities, 
+            model_video
+        ]
     )
 
 if __name__ == "__main__":