Update app.py

app.py

@@ -1,84 +1,104 @@
 import gradio as gr
-import torch
-from sahi.prediction import ObjectPrediction
-from sahi.utils.cv import visualize_object_predictions, read_image
-from ultralyticsplus import YOLO
+import cv2
+import requests
+import os
 
-# Images
-#torch.hub.download_url_to_file('https://ibb.co/zrFWRhJ][img]https://i.ibb.co/37v5Nyz/01-missing-hole-02-jpg-rf-3bd05d8784b12c66da03831d326edd2b.jpg', 'hole.jpg')
-#torch.hub.download_url_to_file('https://ibb.co/qgPMXvN][img]https://i.ibb.co/rp1GTVx/01-mouse-bite-04-jpg-rf-8b3b06d80ac553a92606bdac3e394bb9.jpg', 'mouse.jpg')
-#torch.hub.download_url_to_file('https://ibb.co/0Q35SLN][img]https://i.ibb.co/DLndyFX/12-spur-05-jpg-rf-39e4c0bbc17325d83c938d49e5efdf47.jpg', 'spur.jpg')
+from ultralytics import YOLO
 
-def yolov8_inference(
-    image: gr.inputs.Image = None,
-    model_path: gr.inputs.Dropdown = None,
-    image_size: gr.inputs.Slider = 640,
-    conf_threshold: gr.inputs.Slider = 0.25,
-    iou_threshold: gr.inputs.Slider = 0.45,
-):
-    """
-    YOLOv8 inference function
-    Args:
-        image: Input image
-        model_path: Path to the model
-        image_size: Image size
-        conf_threshold: Confidence threshold
-        iou_threshold: IOU threshold
-    Returns:
-        Rendered image
-    """
-    model = YOLO(model_path)
-    model.conf = conf_threshold
-    model.iou = iou_threshold
-    results = model.predict(image, imgsz=image_size, return_outputs=True)
-    object_prediction_list = []
-    for _, image_results in enumerate(results):
-        if len(image_results)!=0:
-            image_predictions_in_xyxy_format = image_results['det']
-            for pred in image_predictions_in_xyxy_format:
-                x1, y1, x2, y2 = (
-                    int(pred[0]),
-                    int(pred[1]),
-                    int(pred[2]),
-                    int(pred[3]),
-                )
-                bbox = [x1, y1, x2, y2]
-                score = pred[4]
-                category_name = model.model.names[int(pred[5])]
-                category_id = pred[5]
-                object_prediction = ObjectPrediction(
-                    bbox=bbox,
-                    category_id=int(category_id),
-                    score=score,
-                    category_name=category_name,
-                )
-                object_prediction_list.append(object_prediction)
+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)
 
-    image = read_image(image)
-    output_image = visualize_object_predictions(image=image, object_prediction_list=object_prediction_list)
-    return output_image['image']
-        
+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"
+        )
 
-inputs = [
-    gr.inputs.Image(type="filepath", label="Input Image"),
-    gr.inputs.Dropdown(["ethanrom/pcb_def"], 
-                       default="ethanrom/pcb_def", label="Model"),
-    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
+model = YOLO('yolov8_pcb.pt')
+path  = [['spur.jpg'], ['mouse.jpg']]
+video_path = [['video.mp4']]
+
+def show_preds_image(image_path):
+    image = cv2.imread(image_path)
+    outputs = model.predict(source=image_path)
+    results = outputs[0].cpu().numpy()
+    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
+        )
+    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="PCB Defect Detector",
+    examples=path,
+    cache_examples=False,
+)
+
+def show_preds_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    while(cap.isOpened()):
+        ret, frame = cap.read()
+        if ret:
+            frame_copy = frame.copy()
+            outputs = model.predict(source=frame)
+            results = outputs[0].cpu().numpy()
+            for i, det in enumerate(results.boxes.xyxy):
+                cv2.rectangle(
+                    frame_copy,
+                    (int(det[0]), int(det[1])),
+                    (int(det[2]), int(det[3])),
+                    color=(0, 0, 255),
+                    thickness=2,
+                    lineType=cv2.LINE_AA
+                )
+            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
 
-outputs = gr.outputs.Image(type="filepath", label="Output Image")
-title = "YOLOv8 PCB defect detection"
+inputs_video = [
+    gr.components.Video(type="filepath", label="Input Video"),
 
-#examples = [['highway.jpg', 'ethanrom/pcb_def', 640, 0.25, 0.45]]
-demo_app = gr.Interface(
-    fn=yolov8_inference,
-    inputs=inputs,
-    outputs=outputs,
-    title=title,
-    #examples=examples,
-    #cache_examples=True,
-    theme='huggingface',
+]
+outputs_video = [
+    gr.components.Image(type="numpy", label="Output Image"),
+]
+interface_video = gr.Interface(
+    fn=show_preds_video,
+    inputs=inputs_video,
+    outputs=outputs_video,
+    title="PCB Defect Detector",
+    examples=video_path,
+    cache_examples=False,
 )
-demo_app.launch(debug=True, enable_queue=True)
+
+gr.TabbedInterface(
+    [interface_image, interface_video],
+    tab_names=['Image inference', 'Video inference']
+).queue().launch()