Update app.py

app.py

@@ -1,142 +1,185 @@
+import os
 import gradio as gr
-import mlflow
 from ultralytics import YOLO
-from PIL import Image
+import tempfile
 import cv2
 import numpy as np
-import os
+import mlflow
+from datetime import datetime
 import time
-import tempfile
-import sys
-import subprocess
-
-# ==============================
-# MLflow / DagsHub Configuration
-# ==============================
-os.environ["MLFLOW_TRACKING_URI"] = os.getenv("MLFLOW_TRACKING_URI")
-os.environ["MLFLOW_TRACKING_USERNAME"] = os.getenv("MLFLOW_TRACKING_USERNAME")
-os.environ["MLFLOW_TRACKING_PASSWORD"] = os.getenv("MLFLOW_TRACKING_PASSWORD")
-
-dagshub.init(
-    repo_owner="Mosensei7",
-    repo_name="AutonomousVehiclesDetectionDEPI",
-    mlflow=True
-)
-
-mlflow.set_experiment("YOLOv12_Inference")
-
-# ==============================
-# Load YOLOv12 Model
-# ==============================
-model = YOLO("best.pt")  # YOLOv12s weights
-
-# ==============================
-# Inference Logic
-# ==============================
-def run_inference(media_file, media_type):
-    media_path = media_file.name
-
-    with mlflow.start_run(run_name=f"Inference_{int(time.time())}") as run:
-        mlflow.log_param("media_type", media_type)
-        mlflow.log_param("model", "YOLOv12s")
-
-        if media_type == "Image":
-            img = Image.open(media_path).convert("RGB")
-
-            results = model(np.array(img))[0]
-            annotated = results.plot()
-            output_img = Image.fromarray(annotated)
-
-            # Save temp artifacts
-            with tempfile.TemporaryDirectory() as tmp:
-                in_path = os.path.join(tmp, "input.jpg")
-                out_path = os.path.join(tmp, "output.jpg")
-
-                img.save(in_path)
-                output_img.save(out_path)
-
-                mlflow.log_artifact(in_path, "inputs")
-                mlflow.log_artifact(out_path, "outputs")
-
-            mlflow.log_metric("detections", len(results.boxes))
-
-            return output_img, None, run.info.run_id
-
-        else:
-            cap = cv2.VideoCapture(media_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_path = "annotated_output.mp4"
-            writer = cv2.VideoWriter(
-                out_path,
-                cv2.VideoWriter_fourcc(*"mp4v"),
-                fps,
-                (w, h)
-            )
-
-            frame_count = 0
-            total_detections = 0
-
-            while cap.isOpened():
-                ret, frame = cap.read()
-                if not ret:
-                    break
-
-                results = model(frame)[0]
-                annotated = results.plot()
-
-                writer.write(annotated)
-                frame_count += 1
-                total_detections += len(results.boxes)
-
-            cap.release()
-            writer.release()
-
-            mlflow.log_artifact(media_path, "inputs")
-            mlflow.log_artifact(out_path, "outputs")
-            mlflow.log_metric("frames", frame_count)
-            mlflow.log_metric("total_detections", total_detections)
-
-            return None, out_path, run.info.run_id
-
-# ==============================
-# Futuristic UI
-# ==============================
-css = """
-body {
-    background: linear-gradient(135deg, #0f0c29, #302b63, #24243e);
-    color: white;
-    font-family: 'Orbitron', sans-serif;
-}
-.gradio-container {
-    border: 2px solid cyan;
-    border-radius: 20px;
-    box-shadow: 0 0 20px cyan;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("""
-    <h1 style='text-align:center;color:cyan;'>YOLOv12 Autonomous Vehicle Detection</h1>
-    <p style='text-align:center;'>All inferences are logged to DagsHub MLflow</p>
-    """)
-
+import pandas as pd
+from collections import defaultdict
+
+
+DAGSHUB_REPO_OWNER = os.getenv("DAGSHUB_REPO_OWNER")
+DAGSHUB_REPO_NAME = os.getenv("DAGSHUB_REPO_NAME")
+
+MLFLOW_ENABLED = False
+if DAGSHUB_REPO_OWNER and DAGSHUB_REPO_NAME:
+    mlflow.set_tracking_uri(f"https://dagshub.com/{DAGSHUB_REPO_OWNER}/{DAGSHUB_REPO_NAME}.mlflow" )
+    MLFLOW_ENABLED = True
+    print("MLflow tracking is configured for DagsHub.")
+else:
+    print("DagsHub secrets not found. MLflow logging will be disabled.")
+
+os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")
+
+
+MODEL_PATH = "best.pt"
+model = YOLO(MODEL_PATH)
+
+
+def log_image_prediction(input_img_pil, output_image_path, conf, inference_time, detections_df):
+    if not MLFLOW_ENABLED: return
+    try:
+        with mlflow.start_run(run_name=f"Image_Prediction_{datetime.now().strftime('%Y%m%d-%H%M%S')}"):
+            mlflow.log_param("confidence_threshold", conf)
+            mlflow.log_param("prediction_type", "image")
+            mlflow.log_metric("inference_time_seconds", inference_time)
+            mlflow.log_metric("total_detections", len(detections_df))
+
+            if not detections_df.empty:
+                class_counts = defaultdict(int)
+                for _, row in detections_df.iterrows():
+                    class_name = row['class_name']
+                    confidence = row['confidence']
+                    class_counts[class_name] += 1
+                    metric_name = f"detection_{class_name}_{class_counts[class_name]}"
+                    mlflow.log_metric(metric_name, confidence)
+
+            input_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+            input_img_pil.save(input_path)
+            mlflow.log_artifact(input_path, "input_image")
+            mlflow.log_artifact(output_image_path, "output_image")
+            print(f"Successfully logged image prediction.")
+    except Exception as e:
+        print(f"Error logging to MLflow: {e}")
+
+def log_video_prediction(input_path, output_path, conf):
+    if not MLFLOW_ENABLED: return
+    try:
+        with mlflow.start_run(run_name=f"Video_Prediction_{datetime.now().strftime('%Y%m%d-%H%M%S')}"):
+            mlflow.log_param("confidence_threshold", conf)
+            mlflow.log_param("prediction_type", "video")
+            mlflow.log_artifact(input_path, "input")
+            mlflow.log_artifact(output_path, "output")
+            print(f"Successfully logged video prediction.")
+    except Exception as e:
+        print(f"Error logging to MLflow: {e}")
+
+
+def run_image_inference(img_pil, conf=0.25):
+    if img_pil is None: return None, 0.0, pd.DataFrame()
+    img_np = np.array(img_pil)
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    start_time = time.time()
+
+    results = model(img_bgr, conf=conf, iou=0.4, verbose=False, imgsz=640)
+    end_time = time.time()
+    inference_time = end_time - start_time
+    result = results[0]
+    detections = []
+    for box in result.boxes:
+        class_id = int(box.cls.cpu().item())
+        class_name = result.names[class_id]
+        confidence = float(box.conf.cpu().item())
+        detections.append({"class_name": class_name, "confidence": round(confidence, 4)})
+    detections_df = pd.DataFrame(detections)
+    annotated_img = result.plot()
+    annotated_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
+    out_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+    cv2.imwrite(out_path, cv2.cvtColor(annotated_rgb, cv2.COLOR_RGB2BGR))
+    return out_path, inference_time, detections_df
+
+def run_video_inference(video_path, conf=0.25, frame_skip=2):
+    if video_path is None: return None
+    temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+
+    results_generator = model(video_path, conf=conf, iou=0.4, verbose=False, stream=True, imgsz=640)
+    
+    try:
+        first_result = next(results_generator)
+    except StopIteration:
+        return None 
+
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+    h, w = first_result.orig_shape
+
+    output_fps = fps / (frame_skip + 1) if frame_skip > -1 else fps
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(temp_out, fourcc, output_fps, (w, h))
+    
+    out.write(first_result.plot())
+    
+    frame_count = 0
+    for result in results_generator:
+        frame_count += 1
+        if frame_skip > -1 and frame_count % (frame_skip + 1) != 0:
+            continue
+        annotated_frame = result.plot()
+        out.write(annotated_frame)
+        
+    cap.release()
+    out.release()
+    return temp_out
+
+dark_css = """<style> body { background-color: #0f1724; color: #e6eef8; } .gradio-container { background-color: transparent !important; } h1 { color: #ffcc00; } .subtle { color: #9fb0c8; } .card-like { background: rgba(255,255,255,0.03); border-radius: 12px; padding: 12px; } </style>"""
+
+with gr.Blocks() as demo:
+    gr.HTML(dark_css)
+    gr.Markdown("# 🎯 YOLO Detection Studio — Image & Video")
+    gr.Markdown("<div class='subtle'>Upload an image or video, then press Detect.</div>")
     with gr.Row():
-        media = gr.File(label="Upload Image / Video")
-        media_type = gr.Radio(["Image", "Video"], value="Image")
-
-    detect = gr.Button("Run Detection")
-
-    img_out = gr.Image(label="Image Result")
-    vid_out = gr.Video(label="Video Result")
-    run_id = gr.Textbox(label="MLflow Run ID")
-
-    detect.click(
-        run_inference,
-        inputs=[media, media_type],
-        outputs=[img_out, vid_out, run_id]
-    )
-
-demo.launch(share=True)
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.TabItem("Image"):
+                    image_input = gr.Image(type="pil", label="Upload Image")
+                    img_conf = gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="Confidence Threshold")
+                    img_detect_btn = gr.Button("🔍 Detect Image")
+                    image_output = gr.Image(label="Detected Image")
+                with gr.TabItem("Video"):
+                    video_input = gr.Video(label="Upload Video")
+                    vid_conf = gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="Confidence Threshold")
+                    frame_skip_slider = gr.Slider(-1, 10, value=2, step=1, label="Frame Skip", info="Process 1 frame every (N+1) frames. -1 to process all frames.")
+                    vid_detect_btn = gr.Button("🎬 Detect Video")
+                    video_output = gr.Video(label="Detected Video")
+        with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Options & Status")
+            status = gr.Textbox(label="Status", value="Ready", interactive=False)
+            clear_btn = gr.Button("🧹 Clear Outputs")
+
+    def on_detect_image(img, conf):
+        try:
+            out_path, inference_time, detections_df = run_image_inference(img, conf=conf)
+            log_image_prediction(img, out_path, conf, inference_time, detections_df)
+            status_msg = f"Done. Inference: {inference_time:.2f}s. Detections: {len(detections_df)}."
+            if MLFLOW_ENABLED: status_msg += " Logged to DagsHub."
+            return out_path, status_msg
+        except Exception as e:
+            return None, f"Error: {e}"
+
+    def on_detect_video(video_path, conf, frame_skip):
+        try:
+            start_time = time.time()
+            out_path = run_video_inference(video_path, conf=conf, frame_skip=frame_skip)
+            end_time = time.time()
+            if out_path:
+                log_video_prediction(video_path, out_path, conf)
+                status_msg = f"Done — video processed in {end_time - start_time:.2f}s."
+                if MLFLOW_ENABLED: status_msg += " Logged to DagsHub."
+                return out_path, status_msg
+            else:
+                return None, "Could not process video."
+        except Exception as e:
+            import traceback
+            print(traceback.format_exc())
+            return None, f"Error: {e}"
+
+    img_detect_btn.click(fn=on_detect_image, inputs=[image_input, img_conf], outputs=[image_output, status])
+    vid_detect_btn.click(fn=on_detect_video, inputs=[video_input, vid_conf, frame_skip_slider], outputs=[video_output, status])
+    
+    def on_clear(): return None, "Ready", None
+    clear_btn.click(fn=on_clear, inputs=None, outputs=[image_output, status, video_output])
+    
+    demo.launch(server_name="0.0.0.0", share=False)