Update app.py

app.py

@@ -1,185 +1,141 @@
-import os
 import gradio as gr
+import mlflow
+import dagshub
 from ultralytics import YOLO
-import tempfile
+from PIL import Image
 import cv2
 import numpy as np
-import mlflow
-from datetime import datetime
+import os
 import time
-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>")
+import tempfile
+
+# ==============================
+# 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>
+    """)
+
     with gr.Row():
-        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)
+        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)