Update app.py

app.py

@@ -3,109 +3,276 @@ import torch
 import gradio as gr
 import numpy as np
 import os
+import json
+import logging
 import matplotlib.pyplot as plt
-from ultralytics import YOLO, __version__ as ultralytics_version
-import uuid
+from datetime import datetime
+from collections import Counter
+from typing import List, Dict, Any, Optional
+from ultralytics import YOLO
+import ultralytics
+import time
+
+# Set YOLO config directory
+os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
+
+# Set up logging
+logging.basicConfig(
+    filename="app.log",
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+# Directories
+CAPTURED_FRAMES_DIR = "captured_frames"
+OUTPUT_DIR = "outputs"
+os.makedirs(CAPTURED_FRAMES_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+os.chmod(CAPTURED_FRAMES_DIR, 0o777)
+os.chmod(OUTPUT_DIR, 0o777)
+
+# Global variables
+log_entries: List[str] = []
+detected_counts: List[int] = []
+detected_issues: List[str] = []
+gps_coordinates: List[List[float]] = []
+last_metrics: Dict[str, Any] = {}
+frame_count: int = 0
+SAVE_IMAGE_INTERVAL = 1  # Save every frame with detections
 
 # Debug: Check environment
 print(f"Torch version: {torch.__version__}")
 print(f"Gradio version: {gr.__version__}")
-print(f"Ultralytics version: {ultralytics_version}")
+print(f"Ultralytics version: {ultralytics.__version__}")
 print(f"CUDA available: {torch.cuda.is_available()}")
 
-# Load YOLOv8 model
+# Load custom YOLO model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 model = YOLO('./data/best.pt').to(device)
+if device == "cuda":
+    model.half()  # Use half-precision (FP16)
+print(f"Model classes: {model.names}")
+
+# Mock service functions
+def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
+    map_path = "map_temp.png"
+    plt.figure(figsize=(4, 4))
+    plt.scatter([x[1] for x in gps_coords], [x[0] for x in gps_coords], c='blue', label='GPS Points')
+    plt.title("Issue Locations Map")
+    plt.xlabel("Longitude")
+    plt.ylabel("Latitude")
+    plt.legend()
+    plt.savefig(map_path)
+    plt.close()
+    return map_path
+
+def send_to_salesforce(data: Dict[str, Any]) -> None:
+    pass  # Minimal mock
+
+def update_metrics(detections: List[Dict[str, Any]]) -> Dict[str, Any]:
+    counts = Counter([det["label"] for det in detections])
+    return {
+        "items": [{"type": k, "count": v} for k, v in counts.items()],
+        "total_detections": len(detections),
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    }
+
+def generate_line_chart() -> Optional[str]:
+    if not detected_counts:
+        return None
+    plt.figure(figsize=(4, 2))
+    plt.plot(detected_counts[-50:], marker='o', color='#FF8C00')
+    plt.title("Detections Over Time")
+    plt.xlabel("Frame")
+    plt.ylabel("Count")
+    plt.grid(True)
+    plt.tight_layout()
+    chart_path = "chart_temp.png"
+    plt.savefig(chart_path)
+    plt.close()
+    return chart_path
+
+def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
+    global frame_count, last_metrics, detected_counts, detected_issues, gps_coordinates, log_entries
+    frame_count = 0
+    detected_counts.clear()
+    detected_issues.clear()
+    gps_coordinates.clear()
+    log_entries.clear()
+    last_metrics = {}
 
-def process_video(video, output_folder="detected_frames", plot_graphs=False):
     if video is None:
-        yield "Error: No video uploaded", []
-        return
-    
-    # Create output folder if it doesn't exist
-    if not os.path.exists(output_folder):
-        os.makedirs(output_folder)
-    
+        log_entries.append("Error: No video uploaded")
+        logging.error("No video uploaded")
+        return "processed_output.mp4", json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None
+
+    start_time = time.time()
     cap = cv2.VideoCapture(video)
     if not cap.isOpened():
-        yield "Error: Could not open video file", []
-        return
-    
-    frame_width, frame_height = 320, 240  # Smaller resolution
-    frame_count = 0
-    frame_skip = 5  # Process every 5th frame
-    max_frames = 100  # Limit for testing
-    confidence_scores = []  # Store confidence scores for plotting
-    detected_frame_paths = []  # Store paths of frames with detections
-    
+        log_entries.append("Error: Could not open video file")
+        logging.error("Could not open video file")
+        return "processed_output.mp4", json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None
+
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    expected_duration = total_frames / fps if fps > 0 else 0
+    log_entries.append(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
+    logging.info(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
+    print(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
+
+    out_width, out_height = resize_width, resize_height
+    output_path = "processed_output.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'H264')  # Use H264 codec
+    out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height))
+    if not out.isOpened():
+        log_entries.append("Error: Failed to initialize video writer")
+        logging.error("Failed to initialize video writer")
+        cap.release()
+        return "processed_output.mp4", json.dumps({"error": "Failed to initialize video writer"}, indent=2), "\n".join(log_entries), [], None, None
+
+    processed_frames = 0
+    all_detections = []
+    frame_times = []
+    detection_frame_count = 0
+    output_frame_count = 0
+
     while True:
         ret, frame = cap.read()
-        if not ret or frame_count > max_frames:
+        if not ret:
             break
-            
         frame_count += 1
         if frame_count % frame_skip != 0:
             continue
-            
-        frame = cv2.resize(frame, (frame_width, frame_height))
-        print(f"Processing frame {frame_count}")
-        
-        # Run YOLOv8 inference
-        results = model(frame)
-        
-        # Save and yield frame if objects are detected
-        if results[0].boxes is not None and len(results[0].boxes) > 0:
-            annotated_frame = results[0].plot()
-            frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.jpg")
-            cv2.imwrite(frame_filename, annotated_frame)
-            detected_frame_paths.append(frame_filename)
-            
-            # Collect confidence scores for plotting
-            confs = results[0].boxes.conf.cpu().numpy()
-            confidence_scores.extend(confs)
-            
-            # Yield current status and gallery
-            yield f"Processed frame {frame_count} with detections", detected_frame_paths[:]
-    
+        processed_frames += 1
+        frame_start = time.time()
+
+        frame = cv2.resize(frame, (out_width, out_height))
+        results = model(frame, verbose=False, conf=0.5, iou=0.7)
+        annotated_frame = results[0].plot()
+
+        frame_detections = []
+        for detection in results[0].boxes:
+            cls = int(detection.cls)
+            conf = float(detection.conf)
+            box = detection.xyxy[0].cpu().numpy().astype(int).tolist()
+            label = model.names[cls]
+            frame_detections.append({"label": label, "box": box, "conf": conf})
+            log_entries.append(f"Frame {frame_count}: Detected {label} with confidence {conf:.2f}")
+            logging.info(f"Frame {frame_count}: Detected {label} with confidence {conf:.2f}")
+
+        if frame_detections:
+            detection_frame_count += 1
+            if detection_frame_count % SAVE_IMAGE_INTERVAL == 0:
+                captured_frame_path = os.path.join(CAPTURED_FRAMES_DIR, f"detected_{frame_count}.jpg")
+                if not cv2.imwrite(captured_frame_path, annotated_frame):
+                    log_entries.append(f"Error: Failed to save {captured_frame_path}")
+                    logging.error(f"Failed to save {captured_frame_path}")
+                else:
+                    detected_issues.append(captured_frame_path)
+                    if len(detected_issues) > 100:
+                        detected_issues.pop(0)
+
+        # Write frame and duplicates
+        out.write(annotated_frame)
+        output_frame_count += 1
+        if frame_skip > 1:
+            for _ in range(frame_skip - 1):
+                out.write(annotated_frame)
+                output_frame_count += 1
+
+        detected_counts.append(len(frame_detections))
+        gps_coord = [17.385044 + (frame_count * 0.0001), 78.486671 + (frame_count * 0.0001)]
+        gps_coordinates.append(gps_coord)
+        for det in frame_detections:
+            det["gps"] = gps_coord
+        all_detections.extend(frame_detections)
+
+        frame_time = (time.time() - frame_start) * 1000
+        frame_times.append(frame_time)
+        detection_summary = {
+            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "frame": frame_count,
+            "cracks": sum(1 for det in frame_detections if det["label"] == "crack"),
+            "potholes": sum(1 for det in frame_detections if det["label"] == "pothole"),
+            "gps": gps_coord,
+            "processing_time_ms": frame_time
+        }
+        log_entries.append(json.dumps(detection_summary, indent=2))
+        if len(log_entries) > 50:
+            log_entries.pop(0)
+
+    # Pad remaining frames if needed
+    while output_frame_count < total_frames and annotated_frame is not None:
+        out.write(annotated_frame)
+        output_frame_count += 1
+
+    last_metrics = update_metrics(all_detections)
+    send_to_salesforce({
+        "detections": all_detections,
+        "metrics": last_metrics,
+        "timestamp": detection_summary["timestamp"] if all_detections else datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "frame_count": frame_count,
+        "gps_coordinates": gps_coordinates[-1] if gps_coordinates else [0, 0]
+    })
+
+    cap.release()
+    out.release()
+
+    cap = cv2.VideoCapture(output_path)
+    output_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    output_fps = cap.get(cv2.CAP_PROP_FPS)
+    output_duration = output_frames / output_fps if output_fps > 0 else 0
     cap.release()
-    
-    # Generate confidence score plot if requested
-    if plot_graphs and confidence_scores:
-        plt.figure(figsize=(10, 5))
-        plt.hist(confidence_scores, bins=20, color='blue', alpha=0.7)
-        plt.title('Distribution of Confidence Scores')
-        plt.xlabel('Confidence Score')
-        plt.ylabel('Frequency')
-        graph_path = os.path.join(output_folder, "confidence_histogram.png")
-        plt.savefig(graph_path)
-        plt.close()
-        detected_frame_paths.append(graph_path)
-    
-    # Final yield with all results
-    status = f"Saved {len(detected_frame_paths)} frames with detections in {output_folder}. {f'Graph saved as {graph_path}' if plot_graphs and confidence_scores else ''}"
-    yield status, detected_frame_paths
+
+    total_time = time.time() - start_time
+    avg_frame_time = sum(frame_times) / len(frame_times) if frame_times else 0
+    log_entries.append(f"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
+    log_entries.append(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms, Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
+    logging.info(f"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
+    logging.info(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms, Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
+    print(f"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
+    print(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms, Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
+
+    chart_path = generate_line_chart()
+    map_path = generate_map(gps_coordinates[-5:], all_detections)
+
+    return (
+        output_path,
+        json.dumps(last_metrics, indent=2),
+        "\n".join(log_entries[-10:]),
+        detected_issues,
+        chart_path,
+        map_path
+    )
 
 # Gradio interface
-with gr.Blocks() as iface:
-    gr.Markdown("# YOLOv8 Object Detection - Real-time Frame Output")
-    gr.Markdown("Upload a short video to view frames with detections immediately in a gallery. Optionally generate a confidence score graph.")
-    
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
+    gr.Markdown("# Crack and Pothole Detection Dashboard")
     with gr.Row():
-        video_input = gr.Video(label="Upload Video")
-        output_folder = gr.Textbox(label="Output Folder", value="detected_frames")
-        plot_graphs = gr.Checkbox(label="Generate Confidence Score Graph", value=False)
-    
-    submit_button = gr.Button("Process Video")
-    
-    status_output = gr.Text(label="Status")
-    gallery_output = gr.Gallery(label="Detected Frames and Graph", preview=True, columns=3)
-    
-    submit_button.click(
-        fn=process_video,
-        inputs=[video_input, output_folder, plot_graphs],
-        outputs=[status_output, gallery_output],
-        concurrency_limit=1
+        with gr.Column(scale=3):
+            video_input = gr.Video(label="Upload Video")
+            width_slider = gr.Slider(320, 640, value=320, label="Output Width", step=1)
+            height_slider = gr.Slider(240, 480, value=240, label="Output Height", step=1)
+            skip_slider = gr.Slider(1, 10, value=5, label="Frame Skip", step=1)
+            process_btn = gr.Button("Process Video", variant="primary")
+        with gr.Column(scale=1):
+            metrics_output = gr.Textbox(label="Detection Metrics", lines=5, interactive=False)
+    with gr.Row():
+        video_output = gr.Video(label="Processed Video")
+        issue_gallery = gr.Gallery(label="Detected Issues", columns=4, height="auto", object_fit="contain")
+    with gr.Row():
+        chart_output = gr.Image(label="Detection Trend")
+        map_output = gr.Image(label="Issue Locations Map")
+    with gr.Row():
+        logs_output = gr.Textbox(label="Logs", lines=5, interactive=False)
+
+    process_btn.click(
+        process_video,
+        inputs=[video_input, width_slider, height_slider, skip_slider],
+        outputs=[video_output, metrics_output, logs_output, issue_gallery, chart_output, map_output]
     )
 
 if __name__ == "__main__":