app.py

import gradio as gr
import cv2
import requests
import os
import numpy as np

from ultralytics import YOLO

file_urls = [
    'https://www.dropbox.com/s/b5g97xo901zb3ds/pothole_example.jpg?dl=1',
    'https://www.dropbox.com/s/86uxlxxlm1iaexa/pothole_screenshot.png?dl=1',
    'https://www.dropbox.com/s/7sjfwncffg8xej2/video_7.mp4?dl=1'
]

def download_file(url, save_name):
    if not os.path.exists(save_name):
        file = requests.get(url)
        open(save_name, 'wb').write(file.content)

for i, url in enumerate(file_urls):
    if 'mp4' in file_urls[i]:
        download_file(file_urls[i], f"video.mp4")
    else:
        download_file(file_urls[i], f"image_{i}.jpg")

model = YOLO('best.pt')
path = [['image_0.jpg'], ['image_1.jpg']]
video_path = [['video.mp4']]

def save_annotation(image_path, results):
    height, width, _ = cv2.imread(image_path).shape
    annotation_txt = ""
    for i, det in enumerate(results.boxes.xyxy):
        # YOLO format: class x_center y_center width height
        class_id = int(results.names[int(det[5])])
        x_center, y_center, bbox_width, bbox_height = det[0], det[1], det[2] - det[0], det[3] - det[1]
        annotation_txt += f"{class_id} {x_center / width:.6f} {y_center / height:.6f} {bbox_width / width:.6f} {bbox_height / height:.6f}\n"
    return annotation_txt

def show_preds_image(image_path):
    image = cv2.imread(image_path)
    outputs = model.predict(source=image_path)
    results = outputs[0].cpu().numpy()
    
    annotation_txt = save_annotation(image_path, results)

    for i, det in enumerate(results.boxes.xyxy):
        cv2.rectangle(
            image,
            (int(det[0]), int(det[1])),
            (int(det[2]), int(det[3])),
            color=(0, 0, 255),
            thickness=2,
            lineType=cv2.LINE_AA
        )
    
    # Save YOLO format annotation to a txt file
    annotation_filename = f"annotation_{os.path.basename(image_path).split('.')[0]}.txt"
    with open(annotation_filename, 'w') as f:
        f.write(annotation_txt)
    
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

inputs_image = [gr.components.Image(type="filepath", label="Input Image"),]
outputs_image = [gr.components.Image(type="numpy", label="Output Image"),]
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Pothole detector",
    examples=path,
    cache_examples=False,
)

interface_image.launch(debug=True)