Update app.py

app.py

@@ -2,41 +2,48 @@ 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):
-    url = url
     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"
-        )
+        download_file(file_urls[i], f"video.mp4")
     else:
-        download_file(
-            file_urls[i],
-            f"image_{i}.jpg"
-        )
+        download_file(file_urls[i], f"image_{i}.jpg")
 
 model = YOLO('best.pt')
-path  = [['image_0.jpg'], ['image_1.jpg']]
+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,
@@ -46,14 +53,16 @@ def show_preds_image(image_path):
             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"),
-]
+
+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,