app.py

import gradio as gr
import numpy as np
import cv2
import os
from ultralytics import YOLO

# Load the YOLO model
model = YOLO('best.pt')

# Function for image processing
def show_preds_image(image_path):
    image = cv2.imread(image_path)
    results = model.predict(source=image_path)
    annotated_image = results[0].plot()
    return cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

# Function for video processing
def show_preds_video(video_path):
    cap = cv2.VideoCapture(video_path)
    out_frames = []
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model.predict(source=frame)
        annotated_frame = results[0].plot()
        out_frames.append(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
    cap.release()

    # Save the annotated video
    output_path = "annotated_video.mp4"
    height, width, _ = out_frames[0].shape
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    for frame in out_frames:
        writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    writer.release()
    return output_path

# Gradio interfaces
inputs_image = gr.Image(type="filepath", label="Input Image")
outputs_image = gr.Image(type="numpy", label="Output Image")
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Taxi & License Plate Detection with Image"    
)

inputs_video = gr.Video(label="Input Video")
outputs_video = gr.Video(label="Annotated Output")
interface_video = gr.Interface(
    fn=show_preds_video,
    inputs=inputs_video,
    outputs=outputs_video,
    title="Taxi & License Plate Detection with Video"
)

gr.TabbedInterface(
    [interface_image, interface_video],
    tab_names=['Image Inference', 'Video Inference']
).launch(share=True)