Update app.py

app.py

@@ -1,38 +1,36 @@
-#!/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 spaces
 import torch
+from datetime import datetime
 from transformers import AutoModel, AutoTokenizer
-from modelscope.hub.snapshot_download import snapshot_download
+import gradio as gr
 from PIL import Image
 from decord import VideoReader, cpu
 import os
 import gc
-import random
-import io
 import tempfile
 from ultralytics import YOLO
 import numpy as np
 import cv2
-import gradio as gr
+from modelscope.hub.snapshot_download import snapshot_download
 
-# Set device for torch computations
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Initialize GPU
+@spaces.GPU
+def initialize_gpu():
+    if torch.cuda.is_available():
+        torch.randn(10).cuda()
+initialize_gpu()
+
+# Load YOLO model
+YOLO_MODEL = YOLO('best_yolov11.pt')  # Keep this file in repo root
+
+# Model configuration
+MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
+model_dir = snapshot_download(MODEL_NAME, cache_dir='./models')
 
-# Load YOLOv11 model (remove device argument because it's not supported)
-YOLO_MODEL = YOLO('best_yolov11.pt')
+# Device setup
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# File type validation
+# File validation
 IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
 VIDEO_EXTENSIONS = {'.mp4', '.mkv', '.mov', '.avi', '.flv', '.wmv', '.webm', '.m4v'}
 
@@ -45,507 +43,271 @@ def is_image(filename):
 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
-
+@spaces.GPU
 def load_model_and_tokenizer():
-    """Load a fresh instance of the model and tokenizer"""
+    """Load 4-bit quantized model for memory efficiency"""
     try:
-        # Clear GPU memory if using CUDA
-        if DEVICE == "cuda":
-            torch.cuda.empty_cache()
-            gc.collect()
+        torch.cuda.empty_cache()
+        gc.collect()
         
         model = AutoModel.from_pretrained(
-            model_path,
+            model_dir,
             attn_implementation='sdpa',
             trust_remote_code=True,
-            torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
-            device_map='auto'
+            load_in_4bit=True,
+            device_map="auto",
+            torch_dtype=torch.bfloat16
         )
 
         tokenizer = AutoTokenizer.from_pretrained(
-            model_path,
+            model_dir,
             trust_remote_code=True
         )
-        model.eval()
         processor = model.init_processor(tokenizer)
+        model.eval()
         return model, tokenizer, processor
     except Exception as e:
-        print(f"Error loading model: {str(e)}")
+        print(f"Model loading error: {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)]
+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"
+    }
 
-    # Split frame indices into chunks
-    chunks = [
-        frame_idx[i:i + MAX_NUM_FRAMES]
-        for i in range(0, len(frame_idx), MAX_NUM_FRAMES)
-    ]
+    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
 
-    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
+    return counts["Worker"], sum(counts.values()) - counts["Worker"], counts
 
+@spaces.GPU
 def detect_people_and_machinery(media_path):
-    """Detect people and machinery using YOLOv11 for both images and videos"""
+    """GPU-accelerated detection"""
     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
-        }
+        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"
+        ]}
 
-        # 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
-
+            sample_rate = max(1, int(fps))
+            
             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
-
+                
+                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:
-            # 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)
-
+            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_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
+            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"Error in YOLO detection: {str(e)}")
+        print(f"Detection error: {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
-
+@spaces.GPU
 def analyze_video_activities(video_path):
-    """Analyze video using mPLUG model with chunking"""
+    """Video analysis with chunk processing"""
     try:
-        all_responses = []
-        chunk_generator = encode_video_in_chunks(video_path)
+        model, tokenizer, processor = load_model_and_tokenizer()
+        responses = []
         
-        for chunk_idx, video_frames in chunk_generator:
-            # Load fresh model instance for each chunk
-            model, tokenizer, processor = load_model_and_tokenizer()
+        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)
             
-            # 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}")
+            response = model.generate(**inputs, max_new_tokens=200)
+            responses.append(response[0])
             
-            # 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]
+        del model, tokenizer, processor
+        torch.cuda.empty_cache()
+        return "\n".join(responses)
+    
     except Exception as e:
-        print(f"Error processing image: {str(e)}")
-        return "Error processing image"
+        print(f"Video analysis error: {str(e)}")
+        return "Activity analysis unavailable"
 
+@spaces.GPU
 def analyze_image_activities(image_path):
-    """Analyze image using mPLUG model"""
+    """Image analysis pipeline"""
     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)
+        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
-        if DEVICE == "cuda":
-            torch.cuda.empty_cache()
-            gc.collect()
-            
-        return response
+        return response[0]
+    
     except Exception as e:
-        print(f"Error analyzing image: {str(e)}")
-        return "Error analyzing image activities"
+        print(f"Image analysis error: {str(e)}")
+        return "Activity analysis unavailable"
 
-# ------------------------------------------------------------------
-# Function to annotate each frame with bounding boxes & counts
-# ------------------------------------------------------------------
+@spaces.GPU
 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
-    """
+    """Video annotation with real-time detection"""
     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()
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    writer = cv2.VideoWriter(annotated_video_path, fourcc, fps, (w, h))
+    temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    writer = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
 
-    while True:
+    while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-
+        
         results = YOLO_MODEL(frame)
-
-        # Dictionary to hold per-frame counts of each class
-        frame_counts = {}
-
+        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]
+            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]
-
-                # Convert coordinates to int
-                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-
+                counts[class_name] = counts.get(class_name, 0) + 1
+                
                 # 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
-        )
-
+                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 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 plus None for video
-        return [day, date, "No media uploaded", "No media uploaded", "No media uploaded", "No media uploaded", None]
+    return temp_file.name
 
+def process_diary(day, date, people, machinery, machinery_types, activities, media):
+    """Main processing pipeline"""
     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)
+        if not media:
+            return [day, date, "No data", "No data", "No data", "No data", None]
+        
+        with tempfile.NamedTemporaryFile(delete=False) as tmp:
+            tmp.write(media.read())
+            media_path = tmp.name
+        
+        detected_people, detected_machinery, machine_types = detect_people_and_machinery(media_path)
+        annotated_video = None
         
-        # 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)
+            activities = analyze_image_activities(media_path)
         else:
-            # If it's a video, do video analysis and 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]
-
+            activities = analyze_video_activities(media_path)
+            annotated_video = annotate_video_with_bboxes(media_path)
+        
+        os.remove(media_path)
+        return [
+            day,
+            date,
+            str(detected_people),
+            str(detected_machinery),
+            ", ".join([f"{k}:{v}" for k,v in machine_types.items()]),
+            activities,
+            annotated_video
+        ]
+    
     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]
+        print(f"Processing error: {str(e)}")
+        return [day, date, "Error", "Error", "Error", "Error", None]
 
-# Create the Gradio interface
-with gr.Blocks(title="Digital Site Diary") as demo:
-    gr.Markdown("# 📝 Digital Site Diary")
+# Gradio Interface
+with gr.Blocks(title="Digital Site Diary", css="video {height: auto !important;}") as demo:
+    gr.Markdown("# 🏗️ Digital Construction 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
+            gr.Markdown("### Site Details")
+            day = gr.Textbox(label="Day Number", value="1")
+            date = gr.Textbox(label="Date", value=datetime.now().strftime("%Y-%m-%d"))
+            media = gr.File(label="Upload Media", file_types=["image", "video"])
+            submit_btn = gr.Button("Generate Report", variant="primary")
+        
         with gr.Column():
-            gr.Markdown("### Model Detection")
+            gr.Markdown("### Safety Report")
             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)
-            # Annotated video output
-            model_annotated_video = gr.Video(label="Annotated Video")
+            model_people = gr.Textbox(label="Worker Count")
+            model_machinery = gr.Textbox(label="Machinery Count")
+            model_machinery_types = gr.Textbox(label="Machinery Breakdown")
+            model_activities = gr.Textbox(label="Activity Analysis", lines=4)
+            model_video = gr.Video(label="Safety Annotations")
 
-    # 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
-        ]
+        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]
     )
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860)