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)