Update app.py

app.py

@@ -6,6 +6,7 @@ import os
 import json
 import logging
 import matplotlib.pyplot as plt
+import zipfile
 from datetime import datetime
 from collections import Counter
 from typing import List, Dict, Any, Optional
@@ -25,10 +26,13 @@ logging.basicConfig(
 
 # Directories
 CAPTURED_FRAMES_DIR = "captured_frames"
+ORIGINAL_FRAMES_DIR = "original_frames"
 OUTPUT_DIR = "outputs"
 os.makedirs(CAPTURED_FRAMES_DIR, exist_ok=True)
+os.makedirs(ORIGINAL_FRAMES_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 os.chmod(CAPTURED_FRAMES_DIR, 0o777)
+os.chmod(ORIGINAL_FRAMES_DIR, 0o777)
 os.chmod(OUTPUT_DIR, 0o777)
 
 # Global variables
@@ -38,23 +42,16 @@ 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
+SAVE_IMAGE_INTERVAL = 1
 
-# Debug: Check environment
-print(f"Torch version: {torch.__version__}")
-print(f"Gradio version: {gr.__version__}")
-print(f"Ultralytics version: {ultralytics.__version__}")
-print(f"CUDA available: {torch.cuda.is_available()}")
-
-# Load custom YOLO model
+# Load 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}")
+    model.half()
+print(f"Using {device}, model classes: {model.names}")
 
-# Mock service functions
+# Helper 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))
@@ -68,7 +65,7 @@ def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) ->
     return map_path
 
 def send_to_salesforce(data: Dict[str, Any]) -> None:
-    pass  # Minimal mock
+    pass  # Placeholder
 
 def update_metrics(detections: List[Dict[str, Any]]) -> Dict[str, Any]:
     counts = Counter([det["label"] for det in detections])
@@ -81,6 +78,7 @@ def update_metrics(detections: List[Dict[str, Any]]) -> Dict[str, Any]:
 def generate_line_chart() -> Optional[str]:
     if not detected_counts:
         return None
+    chart_path = "chart_temp.png"
     plt.figure(figsize=(4, 2))
     plt.plot(detected_counts[-50:], marker='o', color='#FF8C00')
     plt.title("Detections Over Time")
@@ -88,11 +86,20 @@ def generate_line_chart() -> Optional[str]:
     plt.ylabel("Count")
     plt.grid(True)
     plt.tight_layout()
-    chart_path = "chart_temp.png"
     plt.savefig(chart_path)
     plt.close()
     return chart_path
 
+def create_zip_from_directory(dir_path: str, zip_filename: str) -> str:
+    zip_path = os.path.join(OUTPUT_DIR, zip_filename)
+    with zipfile.ZipFile(zip_path, 'w') as zipf:
+        for root, _, files in os.walk(dir_path):
+            for file in files:
+                full_path = os.path.join(root, file)
+                zipf.write(full_path, arcname=file)
+    return zip_path
+
+# Main function
 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
@@ -102,56 +109,30 @@ def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
     log_entries.clear()
     last_metrics = {}
 
+    for dir_ in [CAPTURED_FRAMES_DIR, ORIGINAL_FRAMES_DIR]:
+        for file in os.listdir(dir_):
+            os.remove(os.path.join(dir_, file))
+
     if video is None:
         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
+        return None, json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None, None, None
 
-    start_time = time.time()
     cap = cv2.VideoCapture(video)
     if not cap.isOpened():
         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
+        return None, json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, 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"
-    # Try codecs in order of preference
-    codecs = [('mp4v', '.mp4'), ('MJPG', '.avi'), ('XVID', '.avi')]
-    out = None
-    for codec, ext in codecs:
-        fourcc = cv2.VideoWriter_fourcc(*codec)
-        output_path = f"processed_output{ext}"
-        out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height))
-        if out.isOpened():
-            log_entries.append(f"Using codec: {codec}, output: {output_path}")
-            logging.info(f"Using codec: {codec}, output: {output_path}")
-            break
-        else:
-            log_entries.append(f"Failed to initialize codec: {codec}")
-            logging.warning(f"Failed to initialize codec: {codec}")
-
-    if not out or not out.isOpened():
-        log_entries.append("Error: All codecs failed to initialize video writer")
-        logging.error("All codecs failed to initialize video writer")
-        cap.release()
-        return "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (resize_width, resize_height))
 
-    processed_frames = 0
     all_detections = []
     frame_times = []
     detection_frame_count = 0
-    output_frame_count = 0
-    last_annotated_frame = None
+    start_time = time.time()
 
     while True:
         ret, frame = cap.read()
@@ -160,77 +141,47 @@ def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
         frame_count += 1
         if frame_count % frame_skip != 0:
             continue
-        processed_frames += 1
-        frame_start = time.time()
 
-        frame = cv2.resize(frame, (out_width, out_height))
+        frame = cv2.resize(frame, (resize_width, resize_height))
         results = model(frame, verbose=False, conf=0.5, iou=0.7)
         annotated_frame = results[0].plot()
 
+        # Save original frame
+        original_path = os.path.join(ORIGINAL_FRAMES_DIR, f"frame_{frame_count}.jpg")
+        cv2.imwrite(original_path, frame)
+
         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]
-            if label != 'Crocodile':  # Ignore irrelevant class
-                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}")
+            frame_detections.append({"label": label, "box": box, "conf": conf})
 
         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
-        last_annotated_frame = annotated_frame
-        if frame_skip > 1:
-            for _ in range(frame_skip - 1):
-                out.write(annotated_frame)
-                output_frame_count += 1
+                captured_path = os.path.join(CAPTURED_FRAMES_DIR, f"frame_{frame_count}.jpg")
+                cv2.imwrite(captured_path, annotated_frame)
+                detected_issues.append(captured_path)
+                if len(detected_issues) > 100:
+                    detected_issues.pop(0)
 
-        detected_counts.append(len(frame_detections))
+        out.write(annotated_frame)
         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
+        detected_counts.append(len(frame_detections))
+        frame_time = (time.time() - start_time) * 1000
         frame_times.append(frame_time)
-        detection_summary = {
-            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            "frame": frame_count,
-            "longitudinal": sum(1 for det in frame_detections if det["label"] == "Longitudinal"),
-            "pothole": sum(1 for det in frame_detections if det["label"] == "Pothole"),
-            "transverse": sum(1 for det in frame_detections if det["label"] == "Transverse"),
-            "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
-    while output_frame_count < total_frames and last_annotated_frame is not None:
-        out.write(last_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"),
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
         "frame_count": frame_count,
         "gps_coordinates": gps_coordinates[-1] if gps_coordinates else [0, 0]
     })
@@ -238,23 +189,10 @@ def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
     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()
-
-    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)
+    originals_zip = create_zip_from_directory(ORIGINAL_FRAMES_DIR, "original_images.zip")
+    annotated_zip = create_zip_from_directory(CAPTURED_FRAMES_DIR, "annotated_images.zip")
 
     return (
         output_path,
@@ -262,12 +200,15 @@ def process_video(video, resize_width=320, resize_height=240, frame_skip=5):
         "\n".join(log_entries[-10:]),
         detected_issues,
         chart_path,
-        map_path
+        map_path,
+        originals_zip,
+        annotated_zip
     )
 
-# Gradio interface
+# Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
-    gr.Markdown("# Road Defect Detection Dashboard")
+    gr.Markdown("# Crack and Pothole Detection Dashboard")
+
     with gr.Row():
         with gr.Column(scale=3):
             video_input = gr.Video(label="Upload Video")
@@ -277,20 +218,36 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
             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)
 
+    with gr.Row():
+        originals_zip_out = gr.File(label="Download Original Images (ZIP)")
+        annotated_zip_out = gr.File(label="Download Annotated Images (ZIP)")
+
     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]
+        outputs=[
+            video_output,
+            metrics_output,
+            logs_output,
+            issue_gallery,
+            chart_output,
+            map_output,
+            originals_zip_out,
+            annotated_zip_out
+        ]
     )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()