initial commit

app.py

@@ -0,0 +1,162 @@
+from PIL import Image
+import gradio as gr
+import cv2
+from ultralytics import ASSETS, YOLO
+import tempfile
+import numpy as np
+import time
+
+def load_model(model_name):
+    """Loads the specified YOLO model for either segmentation or detection."""
+    if model_name == "yolov9c-seg":
+        model_path = "yolov9c-seg.pt"
+    elif model_name == "yolov9e-seg":
+        model_path = "yolov9e-seg.pt"
+    elif model_name == "yolov9c":
+        model_path = "yolov9c.pt"
+    elif model_name == "yolov9e":
+        model_path = "yolov9e.pt"
+    else:
+        raise ValueError(f"Invalid model name: {model_name}")
+    
+    return YOLO(model_path)
+
+def predict_image(img, conf_threshold, iou_threshold, task="detection", model_name=None):
+    """Predicts and plots results in an image using YOLO model with adjustable confidence and IOU thresholds."""
+    if task == "segmentation":
+        if not model_name:
+            model_name = "yolov9c-seg"
+        elif model_name not in ["yolov9c-seg", "yolov9e-seg"]:
+            raise ValueError(f"Invalid model name for segmentation: {model_name}")
+    elif task == "detection":
+        if not model_name:
+            model_name = "yolov9c"
+        elif model_name not in ["yolov9c", "yolov9e"]:
+            raise ValueError(f"Invalid model name for detection: {model_name}")
+    else:
+        raise ValueError(f"Invalid task: {task}. Choose either 'segmentation' or 'detection'.")
+
+    model = load_model(model_name)
+    results = model.predict(
+        source=img,
+        conf=conf_threshold,
+        iou=iou_threshold,
+        show_labels=True,
+        show_conf=True,
+        imgsz=640,
+    )
+
+    for r in results:
+        im_array = r.plot()
+        im = Image.fromarray(im_array[..., ::-1])
+    
+    return im
+
+def predict_image_with_task(img, conf_threshold, iou_threshold, task, model_name):
+
+    return predict_image(img, conf_threshold, iou_threshold, task, model_name)
+
+image_iface = gr.Interface(
+    fn=predict_image_with_task,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
+        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
+        gr.Dropdown(choices=["detection", "segmentation"], value="detection", label="Task"),
+        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov9c-seg", "yolov9e-seg"], value="yolov9c", label="Model"),
+    ],
+    outputs=gr.Image(type="pil", label="Result"),
+    title="X509",
+    description="Upload images for inference. Choose task and corresponding model.",
+    examples=[
+        ["cars.jpg", 0.25, 0.45, "detection", "yolov9c"],
+    ],
+)
+
+def predict_video(video_path, conf_threshold, iou_threshold, task="detection", model_name=None):
+    """Predicts and processes video frames using YOLO model with adjustable confidence and IOU thresholds."""
+    if task == "segmentation":
+        if not model_name:
+            model_name = "yolov9c-seg"
+        elif model_name not in ["yolov9c-seg", "yolov9e-seg"]:
+            raise ValueError(f"Invalid model name for segmentation: {model_name}")
+    elif task == "detection":
+        if not model_name:
+            model_name = "yolov9c"
+        elif model_name not in ["yolov9c", "yolov9e"]:
+            raise ValueError(f"Invalid model name for detection: {model_name}")
+    else:
+        raise ValueError(f"Invalid task: {task}. Choose either 'segmentation' or 'detection'.")
+
+    model = load_model(model_name)
+    cap = cv2.VideoCapture(video_path)
+
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    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_video_path = tempfile.mktemp(suffix=".mp4")
+    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
+
+    frame_count = 0
+    start_time = time.time()
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_count += 1
+
+        elapsed_time = time.time() - start_time
+        current_fps = frame_count / elapsed_time
+
+        pil_img = Image.fromarray(frame[..., ::-1])
+        results = model.predict(
+            source=pil_img,
+            conf=conf_threshold,
+            iou=iou_threshold,
+            show_labels=True,
+            show_conf=True,
+            imgsz=640,
+        )
+
+        for r in results:
+            im_array = r.plot()
+            processed_frame = Image.fromarray(im_array[..., ::-1])
+            frame = cv2.cvtColor(np.array(processed_frame), cv2.COLOR_RGB2BGR)
+
+            cv2.putText(frame, f"FPS: {current_fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+            out.write(frame)
+    
+    cap.release()
+    out.release()
+
+    return temp_video_path
+
+def predict_video_with_task(video_path, conf_threshold, iou_threshold, task, model_name):
+
+    return predict_video(video_path, conf_threshold, iou_threshold, task, model_name)
+
+video_iface = gr.Interface(
+    fn=predict_video_with_task,
+    inputs=[
+        gr.Video(label="Upload Video", interactive=True),
+        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
+        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
+        gr.Dropdown(choices=["detection", "segmentation"], value="detection", label="Task"),
+        gr.Dropdown(choices=["yolov9c", "yolov9e", "yolov9c-seg", "yolov9e-seg"], value="yolov9c", label="Model"),
+    ],
+    outputs=gr.File(label="Result"),
+    title="X509",
+    description="Upload video for inference. Choose task and corresponding model.",
+    examples=[
+        ["VID_20240517112011.mp4", 0.25, 0.45, "detection", "yolov9c"],
+    ]
+)
+
+production = gr.TabbedInterface([image_iface, video_iface], ["Image Inference", "Video Inference"])
+
+if __name__ == '__main__':
+    production.launch()