Update app.py

app.py

@@ -1,4 +1,14 @@
-import gradio as gr
+#!/usr/bin/env python
+# encoding: utf-8
+import spaces
+import torch
+
+# GPU initialization using Spaces decorator
+@spaces.GPU
+def debug():
+    torch.randn(10).cuda()
+debug()
+
 from datetime import datetime
 import torch
 from transformers import AutoModel, AutoTokenizer
@@ -14,17 +24,9 @@ from ultralytics import YOLO
 import numpy as np
 import cv2
 
-# Add this after other model configurations
-YOLO_MODEL = YOLO('best_yolov11.pt')  # Load YOLOv11 model
-
-# Check if CUDA is available
+# Load YOLOv11 model and set device (if supported by YOLO, pass the device)
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Initialize GPU if available
-if DEVICE == "cuda":
-    def debug():
-        torch.randn(10).cuda()
-    debug()
+YOLO_MODEL = YOLO('best_yolov11.pt', device=DEVICE)
 
 # File type validation
 IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
@@ -55,6 +57,516 @@ except Exception as e:
 
 MAX_NUM_FRAMES = 32
 
+def load_model_and_tokenizer():
+    """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()
+        
+        model = AutoModel.from_pretrained(
+            model_path,
+            attn_implementation='sdpa',
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
+            device_map='auto'
+        )
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=True
+        )
+        model.eval()
+        processor = model.init_processor(tokenizer)
+        return model, tokenizer, processor
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        raise
+
+def process_video_chunk(video_frames, model, tokenizer, processor, prompt):
+    """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 "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
+    )
+    # Use DEVICE variable so that CPU-only environments aren’t forced to cuda
+    inputs.to(DEVICE)
+    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 in chunks"""
+    vr = VideoReader(video_path, ctx=cpu(0))
+    sample_fps = round(vr.get_avg_fps() / 1)  # 1 FPS
+    frame_idx = [i for i in range(0, len(vr), sample_fps)]
+
+    # 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 for both images and videos"""
+    try:
+        # Initialize counters with maximum values
+        max_people_count = 0
+        max_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
+        }
+
+        # 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))  # 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:
+            # 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):
+    """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
+    }
+
+    # Process detection results
+    for r in results:
+        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",
+                '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"
+            }
+
+            # Count machinery
+            if conf > 0.5:
+                class_lower = class_name.lower()
+                for key, value in machinery_mapping.items():
+                    if key in class_lower:
+                        machine_types[value] += 1
+                        break
+
+    total_machinery = sum(machine_types.values())
+    return people_count, total_machinery, machine_types
+
+def analyze_video_activities(video_path):
+    """Analyze video using mPLUG model with chunking"""
+    try:
+        all_responses = []
+        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()
+            
+            # 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):
+    """Process single image with mPLUG model"""
+    try:
+        image = Image.open(image_path)
+        messages = [{
+            "role": "user",
+            "content": prompt,
+            "images": [image]
+        }]
+
+        model_messages = []
+        images = []
+
+        for msg in messages:
+            content_str = msg["content"]
+            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
+        )
+        # Use the DEVICE variable for transferring inputs
+        inputs.to(DEVICE)
+        inputs.update({
+            'tokenizer': tokenizer,
+            'max_new_tokens': 100,
+            'decode_text': True,
+        })
+
+        response = model.generate(**inputs)
+        return response[0]
+    except Exception as e:
+        print(f"Error processing image: {str(e)}")
+        return "Error processing image"
+
+def analyze_image_activities(image_path):
+    """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."
+        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):
+    """
+    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:
+            boxes = r.boxes
+            for box in boxes:
+                cls_id = int(box.cls[0])
+                conf = float(box.conf[0])
+                if conf < 0.5:
+                    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)
+
+                # 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
+
+        # 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):
+    """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):
+                with open(media.name, 'rb') as f:
+                    temp_file.write(f.read())
+            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]
+
+# Create the Gradio interface
+import gradio as gr
+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')
+            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)
+            machinery_types = gr.Textbox(
+                label="Number of Machinery Per Type",
+                placeholder="e.g., Excavator: 2, Roller: 1",
+                value="Excavator: 2, Roller: 1"
+            )
+            activities = gr.Textbox(
+                label="Activity",
+                placeholder="e.g., 9 AM: Excavation, 10 AM: Concreting",
+                value="9 AM: Excavation, 10 AM: Concreting",
+                lines=3
+            )
+            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")
+            model_day = gr.Textbox(label="Day")
+            model_date = gr.Textbox(label="Date")
+            model_people = gr.Textbox(label="Total Number of People")
+            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  # The new 7th output
+        ]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=False)
+    return get_file_extension(filename) in IMAGE_EXTENSIONS
+
+def is_video(filename):
+    return get_file_extension(filename) in VIDEO_EXTENSIONS
+
+# 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:
+    print(f"Error downloading model: {str(e)}")
+    model_path = os.path.join(MODEL_CACHE_DIR, MODEL_NAME)
+
+MAX_NUM_FRAMES = 32
+
 def load_model_and_tokenizer():
     """Load a fresh instance of the model and tokenizer"""
     try: