Update app.py

app.py

@@ -1,66 +1,39 @@
+#!/usr/bin/env python
 import spaces
 import torch
+@spaces.GPU
+def debug():
+    torch.randn(10).cuda()
+debug()
+
+import gradio as gr
 from datetime import datetime
+import torch
 from transformers import AutoModel, AutoTokenizer
-import gradio as gr
+from modelscope.hub.snapshot_download import snapshot_download
 from PIL import Image
 from decord import VideoReader, cpu
 import os
 import gc
+import io
 import tempfile
 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
+# Load YOLOv11 model (update the path as needed)
+YOLO_MODEL = YOLO('best_yolov11.pt')
 
-# Fix GLIBCXX dependency
-os.environ['LD_LIBRARY_PATH'] = '/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH'
+# Check if CUDA is available
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Initialize GPU
-@spaces.GPU
-def initialize_gpu():
-    if torch.cuda.is_available():
+# Initialize GPU if available (already done by debug() above)
+if DEVICE == "cuda":
+    def init_debug():
         torch.randn(10).cuda()
-initialize_gpu()
-
-# Load YOLO model with error handling
-try:
-    YOLO_MODEL = YOLO('best_yolov11.pt')
-except Exception as e:
-    raise RuntimeError(f"YOLO model loading failed: {str(e)}")
-
-# Model configuration
-MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
-try:
-    model_dir = snapshot_download(MODEL_NAME, 
-                                 cache_dir='./models',
-                                 revision='main')
-except Exception as e:
-    raise RuntimeError(f"Model download failed: {str(e)}")
-
-# Device setup
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    init_debug()
 
-# File validation
+# File type validation
 IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
 VIDEO_EXTENSIONS = {'.mp4', '.mkv', '.mov', '.avi', '.flv', '.wmv', '.webm', '.m4v'}
 
@@ -73,293 +46,402 @@ def is_image(filename):
 def is_video(filename):
     return get_file_extension(filename) in VIDEO_EXTENSIONS
 
-@spaces.GPU
+# Model configuration
+MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
+MODEL_CACHE_DIR = os.getenv('TRANSFORMERS_CACHE', './models')
+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 8-bit quantized model"""
+    """Load a fresh instance of the model and tokenizer"""
     try:
-        torch.cuda.empty_cache()
-        gc.collect()
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
+            gc.collect()
         
         model = AutoModel.from_pretrained(
-            model_dir,
+            model_path,
             attn_implementation='sdpa',
             trust_remote_code=True,
-            load_in_8bit=True,
-            device_map="auto",
-            torch_dtype=torch.float16
+            torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
+            device_map='auto'
         )
-
         tokenizer = AutoTokenizer.from_pretrained(
-            model_dir,
+            model_path,
             trust_remote_code=True
         )
-        processor = model.init_processor(tokenizer)
         model.eval()
+        processor = model.init_processor(tokenizer)
         return model, tokenizer, processor
     except Exception as e:
-        print(f"Model loading error: {str(e)}")
+        print(f"Error loading model: {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"
-    }
-
-    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
+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
+    )
+    inputs.to(DEVICE)
+    inputs.update({
+        'tokenizer': tokenizer,
+        'max_new_tokens': 100,
+        'decode_text': True,
+    })
+    response = model.generate(**inputs)
+    return response[0]
 
-    return counts["Worker"], sum(counts.values()) - counts["Worker"], counts
+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)]
+    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
 
-@spaces.GPU
 def detect_people_and_machinery(media_path):
-    """GPU-accelerated detection"""
+    """Detect people and machinery using YOLOv11 for images and videos"""
     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"
-        ]}
-
+        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
+        }
         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
             while cap.isOpened():
                 ret, frame = cap.read()
                 if not ret:
                     break
-                
-                if cap.get(cv2.CAP_PROP_POS_FRAMES) % sample_rate == 0:
+                if frame_count % sample_rate == 0:
                     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))
-            
+                    people, _, machine_types = process_yolo_results(results)
+                    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)
+            if isinstance(media_path, str):
+                img = cv2.imread(media_path)
+            else:
+                img = 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
-    
+            max_people_count, _, max_machine_types = process_yolo_results(results)
+        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"Detection error: {str(e)}")
+        print(f"Error in YOLO detection: {str(e)}")
         return 0, 0, {}
 
-@spaces.GPU
+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
+    }
+    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]
+            if class_name.lower() == 'worker' and conf > 0.5:
+                people_count += 1
+            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"
+            }
+            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):
-    """Video analysis with chunk processing"""
+    """Analyze video using mPLUG model with chunking"""
     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()))
-        total_frames = len(vr)
-        
-        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)]
-            
-            inputs = processor(
-                [{"role": "user", "content": "Analyze construction activities", "video_frames": frames}],
-                videos=[frames]
-            ).to(DEVICE)
-            
-            response = model.generate(**inputs, max_new_tokens=200)
-            responses.append(response[0])
-            
-            del frames, inputs
+        all_responses = []
+        chunk_generator = encode_video_in_chunks(video_path)
+        for chunk_idx, video_frames in chunk_generator:
+            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.")
+            response = process_video_chunk(video_frames, model, tokenizer, processor, prompt)
+            all_responses.append(f"Time period {chunk_idx + 1}:\n{response}")
+            del model, tokenizer, processor
             torch.cuda.empty_cache()
-        
-        del model, tokenizer, processor
-        return "\n".join(responses)
-    
+            gc.collect()
+        return "\n\n".join(all_responses)
     except Exception as e:
-        print(f"Video analysis error: {str(e)}")
-        return "Activity analysis unavailable"
+        print(f"Error analyzing video: {str(e)}")
+        return "Error analyzing video activities"
 
-@spaces.GPU
-def analyze_image_activities(image_path):
-    """Image analysis pipeline"""
+def process_image(image_path, model, tokenizer, processor, prompt):
+    """Process single image with mPLUG model"""
     try:
-        model, tokenizer, processor = load_model_and_tokenizer()
-        image = Image.open(image_path).convert("RGB")
-        
+        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(
-            [{"role": "user", "content": "Analyze construction site", "images": [image]}],
-            images=[image]
-        ).to(DEVICE)
-        
-        response = model.generate(**inputs, max_new_tokens=200)
-        
-        del model, tokenizer, processor, image, inputs
-        torch.cuda.empty_cache()
+            model_messages,
+            images=images,
+            videos=None
+        )
+        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"Image analysis error: {str(e)}")
-        return "Activity analysis unavailable"
+        print(f"Error processing image: {str(e)}")
+        return "Error processing image"
 
-@spaces.GPU
-def annotate_video_with_bboxes(video_path):
-    """Video annotation with detection overlay"""
+def analyze_image_activities(image_path):
+    """Analyze image using mPLUG model"""
     try:
-        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))
-
-        frame_count = 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            
-            if frame_count % 5 == 0:
-                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
-                        
-                        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)
-                
-                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)
-            frame_count += 1
-        
-        cap.release()
-        writer.release()
-        return temp_file.name
-    
+        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"Video annotation error: {str(e)}")
-        return None
+        print(f"Error analyzing image: {str(e)}")
+        return "Error analyzing image activities"
+
+def annotate_video_with_bboxes(video_path):
+    """
+    Reads the video frame-by-frame, runs YOLO, draws bounding boxes and summaries,
+    then writes the annotated frames into a new video.
+    """
+    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))
+    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)
+        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
+                x1, y1, x2, y2 = box.xyxy[0]
+                class_name = YOLO_MODEL.names[cls_id]
+                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                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)
+                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
+        )
+        writer.write(frame)
+    cap.release()
+    writer.release()
+    return annotated_video_path
 
-def process_diary(day, date, media):
-    """Main processing pipeline"""
+@spaces.GPU
+def process_diary(day, date, total_people, total_machinery, machinery_types, activities, media):
+    """Process the site diary entry and return outputs along with an annotated video (if video)"""
+    if media is None:
+        return [day, date, "No media uploaded", "No media uploaded", "No media uploaded", "No media uploaded", None]
     try:
-        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
-        
-        try:
-            if is_image(media.name):
-                activities = analyze_image_activities(media_path)
+        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:
-                activities = analyze_video_activities(media_path)
-                annotated_video = annotate_video_with_bboxes(media_path)
-        except Exception as e:
-            activities = f"Analysis error: {str(e)}"
-        
-        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
-        ]
-    
+                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)
+        annotated_video_path = None
+        if is_image(media.name):
+            detected_activities = analyze_image_activities(temp_path)
+        else:
+            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 [day, date, str(detected_people), str(detected_machinery), detected_types_str, detected_activities, annotated_video_path]
     except Exception as e:
-        print(f"Processing error: {str(e)}")
-        return [day, date, "Error", "Error", "Error", "Error", None]
+        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
-with gr.Blocks(title="Digital Site Diary", css="video {height: auto !important;}") as demo:
-    gr.Markdown("# 🏗️ Digital Construction Diary")
-    
+with gr.Blocks(title="Digital Site Diary") as demo:
+    gr.Markdown("# 📝 Digital Site Diary")
     with gr.Row():
         with gr.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")
-        
+            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")
         with gr.Column():
-            gr.Markdown("### Safety Report")
+            gr.Markdown("### Model Detection")
             model_day = gr.Textbox(label="Day")
             model_date = gr.Textbox(label="Date")
-            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")
-
+            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)
+            model_annotated_video = gr.Video(label="Annotated Video")
     submit_btn.click(
-        process_diary,
-        inputs=[day, date, media],
+        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, 
             model_machinery_types, 
-            model_activities, 
-            model_video
+            model_activities,
+            model_annotated_video
         ]
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()