Update app.py

app.py

@@ -1,42 +1,32 @@
-#!/usr/bin/env python
-# encoding: utf-8
-
-import spaces
-import torch
-
-# Initialize GPU (this decorator ensures that GPU resources are allocated)
-@spaces.GPU
-def debug():
-    torch.randn(10).cuda()
-debug()
-
-from datetime import datetime
 import gradio as gr
+from datetime import datetime
+import torch
 from transformers import AutoModel, AutoTokenizer
 from modelscope.hub.snapshot_download import snapshot_download
 from PIL import Image
 from decord import VideoReader, cpu
 import os
 import gc
+import random
+import io
 import tempfile
-import cv2
-import numpy as np
 from ultralytics import YOLO
+import numpy as np
+import cv2
 
-# -------------------------------
-# Model and File Configurations
-# -------------------------------
-
-# Load your custom YOLOv11 model (adjust the path as needed)
-YOLO_MODEL = YOLO('/teamspace/studios/this_studio/best_yolov11.pt')
+# Add this after other model configurations
+YOLO_MODEL = YOLO('/teamspace/studios/this_studio/best_yolov11.pt')  # Load YOLOv11 model
 
 # Check if CUDA is available
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Initialize GPU if available
 if DEVICE == "cuda":
-    # Quick GPU warm-up
-    torch.randn(10).cuda()
+    def debug():
+        torch.randn(10).cuda()
+    debug()
 
-# Valid extensions for media files
+# File type validation
 IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
 VIDEO_EXTENSIONS = {'.mp4', '.mkv', '.mov', '.avi', '.flv', '.wmv', '.webm', '.m4v'}
 
@@ -49,13 +39,14 @@ def is_image(filename):
 def is_video(filename):
     return get_file_extension(filename) in VIDEO_EXTENSIONS
 
-# -------------------------------
-# mPLUG-Owl Model Configuration
-# -------------------------------
+# Model configuration
 MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
 MODEL_CACHE_DIR = os.getenv('TRANSFORMERS_CACHE', './models')
+
+# Create cache directory if it doesn't exist
 os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
 
+# Download and cache the model
 try:
     model_path = snapshot_download(MODEL_NAME, cache_dir=MODEL_CACHE_DIR)
 except Exception as e:
@@ -65,8 +56,9 @@ except Exception as e:
 MAX_NUM_FRAMES = 32
 
 def load_model_and_tokenizer():
-    """Load a fresh instance of the model, tokenizer, and processor."""
+    """Load a fresh instance of the model and tokenizer"""
     try:
+        # Clear GPU memory if using CUDA
         if DEVICE == "cuda":
             torch.cuda.empty_cache()
             gc.collect()
@@ -78,7 +70,11 @@ def load_model_and_tokenizer():
             torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
             device_map='auto'
         )
-        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=True
+        )
         model.eval()
         processor = model.init_processor(tokenizer)
         return model, tokenizer, processor
@@ -86,52 +82,70 @@ def load_model_and_tokenizer():
         print(f"Error loading model: {str(e)}")
         raise
 
-# -------------------------------
-# Video & Image Processing Functions
-# -------------------------------
-
 def process_video_chunk(video_frames, model, tokenizer, processor, prompt):
-    """Process a chunk of video frames with the mPLUG model."""
-    messages = [{
-        "role": "user",
-        "content": prompt,
-        "video_frames": video_frames
-    }]
+    """Process a chunk of video frames with mPLUG model"""
+    messages = [
+        {
+            "role": "user",
+            "content": prompt,
+            "video_frames": video_frames
+        }
+    ]
 
     model_messages = []
     videos = []
+
     for msg in messages:
         content_str = msg["content"]
-        if msg.get("video_frames"):
+        if "video_frames" in msg and msg["video_frames"]:
             content_str += "<|video|>"
             videos.append(msg["video_frames"])
-        model_messages.append({"role": msg["role"], "content": content_str})
-    model_messages.append({"role": "assistant", "content": ""})
-    
-    inputs = processor(model_messages, images=None, videos=videos if videos else None)
+        model_messages.append({
+            "role": msg["role"],
+            "content": content_str
+        })
+
+    model_messages.append({
+        "role": "assistant",
+        "content": ""
+    })
+
+    inputs = processor(
+        model_messages,
+        images=None,
+        videos=videos if videos else None
+    )
     inputs.to('cuda')
     inputs.update({
         'tokenizer': tokenizer,
         'max_new_tokens': 100,
         'decode_text': True,
     })
+
     response = model.generate(**inputs)
     return response[0]
 
 def encode_video_in_chunks(video_path):
-    """Extract frames from a video and yield them in chunks."""
+    """Extract frames from a video in chunks"""
     vr = VideoReader(video_path, ctx=cpu(0))
-    sample_fps = round(vr.get_avg_fps() / 1)  # 1 FPS sampling
+    sample_fps = round(vr.get_avg_fps() / 1)  # 1 FPS
     frame_idx = [i for i in range(0, len(vr), sample_fps)]
-    chunks = [frame_idx[i:i + MAX_NUM_FRAMES] for i in range(0, len(frame_idx), MAX_NUM_FRAMES)]
+
+    # Split frame indices into chunks
+    chunks = [
+        frame_idx[i:i + MAX_NUM_FRAMES]
+        for i in range(0, len(frame_idx), MAX_NUM_FRAMES)
+    ]
+
     for chunk_idx, chunk in enumerate(chunks):
         frames = vr.get_batch(chunk).asnumpy()
         frames = [Image.fromarray(v.astype('uint8')) for v in frames]
         yield chunk_idx, frames
 
 def detect_people_and_machinery(media_path):
-    """Detect people and machinery (using YOLOv11) in an image or video."""
+    """Detect people and machinery using YOLOv11 for both images and videos"""
     try:
+        # Initialize counters with maximum values
         max_people_count = 0
         max_machine_types = {
             "Tower Crane": 0,
@@ -148,49 +162,84 @@ def detect_people_and_machinery(media_path):
             "Other Vehicle": 0
         }
 
+        # Check if input is video
         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))
-            frame_count = 0
+            sample_rate = max(1, int(fps))  # Sample 1 frame per second
+            frame_count = 0  # Initialize frame counter
+
             while cap.isOpened():
                 ret, frame = cap.read()
                 if not ret:
                     break
+
+                # Process every nth frame based on sample rate
                 if frame_count % sample_rate == 0:
                     results = YOLO_MODEL(frame)
                     people, _, machine_types = process_yolo_results(results)
+                    
+                    # Update maximum counts
                     max_people_count = max(max_people_count, people)
                     for k, v in machine_types.items():
                         max_machine_types[k] = max(max_machine_types[k], v)
+
                 frame_count += 1
+
             cap.release()
+
         else:
-            img = cv2.imread(media_path) if isinstance(media_path, str) else cv2.cvtColor(np.array(media_path), cv2.COLOR_RGB2BGR)
+            # Handle single image
+            if isinstance(media_path, str):
+                img = cv2.imread(media_path)
+            else:
+                # Handle PIL Image
+                img = cv2.cvtColor(np.array(media_path), cv2.COLOR_RGB2BGR)
+
             results = YOLO_MODEL(img)
             max_people_count, _, max_machine_types = process_yolo_results(results)
+
+        # Filter out machinery types with zero count
         max_machine_types = {k: v for k, v in max_machine_types.items() if v > 0}
         total_machinery_count = sum(max_machine_types.values())
+
         return max_people_count, total_machinery_count, max_machine_types
+
     except Exception as e:
         print(f"Error in YOLO detection: {str(e)}")
         return 0, 0, {}
 
 def process_yolo_results(results):
-    """Count detected workers and machinery from YOLO results."""
+    """Process YOLO detection results and count people and machinery"""
     people_count = 0
     machine_types = {
-        "Tower Crane": 0, "Mobile Crane": 0, "Compactor/Roller": 0, "Bulldozer": 0,
-        "Excavator": 0, "Dump Truck": 0, "Concrete Mixer": 0, "Loader": 0,
-        "Pump Truck": 0, "Pile Driver": 0, "Grader": 0, "Other Vehicle": 0
+        "Tower Crane": 0,
+        "Mobile Crane": 0,
+        "Compactor/Roller": 0,
+        "Bulldozer": 0,
+        "Excavator": 0,
+        "Dump Truck": 0,
+        "Concrete Mixer": 0,
+        "Loader": 0,
+        "Pump Truck": 0,
+        "Pile Driver": 0,
+        "Grader": 0,
+        "Other Vehicle": 0
     }
+
+    # Process detection results
     for r in results:
-        for box in r.boxes:
+        boxes = r.boxes
+        for box in boxes:
             cls = int(box.cls[0])
             conf = float(box.conf[0])
             class_name = YOLO_MODEL.names[cls]
+            
+            # Count people (Worker class)
             if class_name.lower() == 'worker' and conf > 0.5:
                 people_count += 1
+            
+            # Map YOLO classes to machinery types
             machinery_mapping = {
                 'tower_crane': "Tower Crane",
                 'mobile_crane': "Mobile Crane",
@@ -208,33 +257,46 @@ def process_yolo_results(results):
                 'grader': "Grader",
                 'other_vehicle': "Other Vehicle"
             }
+
+            # Count machinery
             if conf > 0.5:
+                class_lower = class_name.lower()
                 for key, value in machinery_mapping.items():
-                    if key in class_name.lower():
+                    if key in class_lower:
                         machine_types[value] += 1
                         break
-    return people_count, sum(machine_types.values()), machine_types
+
+    total_machinery = sum(machine_types.values())
+    return people_count, total_machinery, machine_types
 
 def analyze_video_activities(video_path):
-    """Analyze a video by processing it in chunks with the mPLUG model."""
+    """Analyze video using mPLUG model with chunking"""
     try:
         all_responses = []
-        for chunk_idx, video_frames in encode_video_in_chunks(video_path):
+        chunk_generator = encode_video_in_chunks(video_path)
+        
+        for chunk_idx, video_frames in chunk_generator:
+            # Load fresh model instance for each chunk
             model, tokenizer, processor = load_model_and_tokenizer()
-            prompt = ("Analyze this construction site video chunk and describe the activities happening. "
-                      "Focus on construction activities, machinery usage, and worker actions.")
+            
+            # Process the chunk
+            prompt = "Analyze this construction site video chunk and describe the activities happening. Focus on construction activities, machinery usage, and worker actions."
             response = process_video_chunk(video_frames, model, tokenizer, processor, prompt)
             all_responses.append(f"Time period {chunk_idx + 1}:\n{response}")
+            
+            # Clean up GPU memory
             del model, tokenizer, processor
             torch.cuda.empty_cache()
             gc.collect()
+        
+        # Combine all responses
         return "\n\n".join(all_responses)
     except Exception as e:
         print(f"Error analyzing video: {str(e)}")
         return "Error analyzing video activities"
 
 def process_image(image_path, model, tokenizer, processor, prompt):
-    """Analyze a single image with the mPLUG model."""
+    """Process single image with mPLUG model"""
     try:
         image = Image.open(image_path)
         messages = [{
@@ -242,22 +304,37 @@ def process_image(image_path, model, tokenizer, processor, prompt):
             "content": prompt,
             "images": [image]
         }]
+
         model_messages = []
         images = []
+
         for msg in messages:
             content_str = msg["content"]
-            if msg.get("images"):
+            if "images" in msg and msg["images"]:
                 content_str += "<|image|>"
                 images.extend(msg["images"])
-            model_messages.append({"role": msg["role"], "content": content_str})
-        model_messages.append({"role": "assistant", "content": ""})
-        inputs = processor(model_messages, images=images, videos=None)
+            model_messages.append({
+                "role": msg["role"],
+                "content": content_str
+            })
+
+        model_messages.append({
+            "role": "assistant",
+            "content": ""
+        })
+
+        inputs = processor(
+            model_messages,
+            images=images,
+            videos=None
+        )
         inputs.to('cuda')
         inputs.update({
             'tokenizer': tokenizer,
             'max_new_tokens': 100,
             'decode_text': True,
         })
+
         response = model.generate(**inputs)
         return response[0]
     except Exception as e:
@@ -265,69 +342,120 @@ def process_image(image_path, model, tokenizer, processor, prompt):
         return "Error processing image"
 
 def analyze_image_activities(image_path):
-    """Wrapper to analyze an image using mPLUG."""
+    """Analyze image using mPLUG model"""
     try:
         model, tokenizer, processor = load_model_and_tokenizer()
-        prompt = ("Analyze this construction site image and describe the activities happening. "
-                  "Focus on construction activities, machinery usage, and worker actions.")
+        prompt = "Analyze this construction site image and describe the activities happening. Focus on construction activities, machinery usage, and worker actions."
         response = process_image(image_path, model, tokenizer, processor, prompt)
+        
         del model, tokenizer, processor
         if DEVICE == "cuda":
             torch.cuda.empty_cache()
             gc.collect()
+            
         return response
     except Exception as e:
         print(f"Error analyzing image: {str(e)}")
         return "Error analyzing image activities"
 
+
+# ------------------------------------------------------------------
+# NEW: Function to annotate each frame with bounding boxes & counts
+# ------------------------------------------------------------------
 def annotate_video_with_bboxes(video_path):
-    """Annotate each frame of the video with bounding boxes and a summary of detected classes."""
+    """
+    Reads the entire video frame-by-frame, runs YOLO, draws bounding boxes,
+    writes a per-frame summary of detected classes on the frame, and saves
+    as a new annotated video. Returns: annotated_video_path
+    """
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
     w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Create a temp file for output
     out_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
     annotated_video_path = out_file.name
     out_file.close()
+
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     writer = cv2.VideoWriter(annotated_video_path, fourcc, fps, (w, h))
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
+
         results = YOLO_MODEL(frame)
+
+        # Dictionary to hold per-frame counts of each class
         frame_counts = {}
+
         for r in results:
-            for box in r.boxes:
+            boxes = r.boxes
+            for box in boxes:
                 cls_id = int(box.cls[0])
                 conf = float(box.conf[0])
                 if conf < 0.5:
-                    continue
+                    continue  # Skip low-confidence
+
                 x1, y1, x2, y2 = box.xyxy[0]
                 class_name = YOLO_MODEL.names[cls_id]
+
+                # Convert to int
                 x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+
+                # Draw bounding box
+                color = (0, 255, 0)
+                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+
                 label_text = f"{class_name} {conf:.2f}"
                 cv2.putText(frame, label_text, (x1, y1 - 6),
                             cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
+
+                # Increment per-frame class count
                 frame_counts[class_name] = frame_counts.get(class_name, 0) + 1
-        summary_str = ", ".join(f"{cls_name}: {count}" for cls_name, count in frame_counts.items())
-        cv2.putText(frame, summary_str, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 0), 2)
+
+        # Build a summary line, e.g. "Worker: 2, Excavator: 1, ..."
+        summary_str = ", ".join(f"{cls_name}: {count}" 
+                                for cls_name, count in frame_counts.items())
+
+        # Put the summary text in the top-left
+        cv2.putText(
+            frame,
+            summary_str,
+            (15, 30),  # position
+            cv2.FONT_HERSHEY_SIMPLEX,
+            1.0,
+            (255, 255, 0),
+            2
+        )
+
         writer.write(frame)
+
     cap.release()
     writer.release()
     return annotated_video_path
 
+
+
+# ----------------------------------------------------------------------------
+# Update process_diary function to also return an annotated video if it's video
+# ----------------------------------------------------------------------------
 def process_diary(day, date, total_people, total_machinery, machinery_types, activities, media):
-    """Combine the inputs into a site diary entry and perform detection/analysis."""
+    """Process the site diary entry"""
     if media is None:
+        # Return 6 text outputs as before + None for video
         return [day, date, "No media uploaded", "No media uploaded", "No media uploaded", "No media uploaded", None]
+
     try:
         if not hasattr(media, 'name'):
             raise ValueError("Invalid file upload")
+            
         file_ext = get_file_extension(media.name)
         if not (is_image(media.name) or is_video(media.name)):
             raise ValueError(f"Unsupported file type: {file_ext}")
+
         with tempfile.NamedTemporaryFile(suffix=file_ext, delete=False) as temp_file:
             temp_path = temp_file.name
             if hasattr(media, 'name') and os.path.exists(media.name):
@@ -336,32 +464,41 @@ def process_diary(day, date, total_people, total_machinery, machinery_types, act
             else:
                 file_content = media.read() if hasattr(media, 'read') else media
                 temp_file.write(file_content if isinstance(file_content, bytes) else file_content.read())
+
         detected_people, detected_machinery, detected_machinery_types = detect_people_and_machinery(temp_path)
+        
+        # Default: no annotated video
         annotated_video_path = None
+
         if is_image(media.name):
+            # If it's an image, do normal image analysis
             detected_activities = analyze_image_activities(temp_path)
         else:
+            # If it's a video, do video analysis & also annotate the video
             detected_activities = analyze_video_activities(temp_path)
             annotated_video_path = annotate_video_with_bboxes(temp_path)
+
         if os.path.exists(temp_path):
             os.remove(temp_path)
+
         detected_types_str = ", ".join([f"{k}: {v}" for k, v in detected_machinery_types.items()])
+        # Return 7 outputs (the first 6 as before, plus the annotated video path)
         return [day, date, str(detected_people), str(detected_machinery), detected_types_str, detected_activities, annotated_video_path]
+
     except Exception as e:
         print(f"Error processing media: {str(e)}")
         return [day, date, "Error processing media", "Error processing media", "Error processing media", "Error processing media", None]
 
-# -------------------------------
-# Gradio Interface Setup
-# -------------------------------
 
+# Create the Gradio interface
 with gr.Blocks(title="Digital Site Diary") as demo:
     gr.Markdown("# 📝 Digital Site Diary")
+    
     with gr.Row():
         # User Input Column
         with gr.Column():
             gr.Markdown("### User Input")
-            day = gr.Textbox(label="Day", value='9')
+            day = gr.Textbox(label="Day",value='9')
             date = gr.Textbox(label="Date", placeholder="YYYY-MM-DD", value=datetime.now().strftime("%Y-%m-%d"))
             total_people = gr.Number(label="Total Number of People", precision=0, value=10)
             total_machinery = gr.Number(label="Total Number of Machinery", precision=0, value=3)
@@ -378,6 +515,7 @@ with gr.Blocks(title="Digital Site Diary") as demo:
             )
             media = gr.File(label="Upload Image/Video", file_types=["image", "video"])
             submit_btn = gr.Button("Submit", variant="primary")
+
         # Model Detection Column
         with gr.Column():
             gr.Markdown("### Model Detection")
@@ -387,12 +525,23 @@ with gr.Blocks(title="Digital Site Diary") as demo:
             model_machinery = gr.Textbox(label="Total Number of Machinery")
             model_machinery_types = gr.Textbox(label="Number of Machinery Per Type")
             model_activities = gr.Textbox(label="Activity", lines=5)
+            # NEW: annotated video output
             model_annotated_video = gr.Video(label="Annotated Video")
+
+    # Connect the submit button to the processing function
     submit_btn.click(
         fn=process_diary,
         inputs=[day, date, total_people, total_machinery, machinery_types, activities, media],
-        outputs=[model_day, model_date, model_people, model_machinery, model_machinery_types, model_activities, model_annotated_video]
+        outputs=[
+            model_day, 
+            model_date, 
+            model_people, 
+            model_machinery, 
+            model_machinery_types, 
+            model_activities,
+            model_annotated_video  # The new 7th output
+        ]
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True) this is my code i want to deploy it on hugging face gradio