Update app.py

app.py

@@ -1,86 +1,105 @@
 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, render_result
+import cv2
+import requests
+import os
 
+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.overrides['conf'] = conf_threshold
-    model.overrides['iou']= iou_threshold
-    model.overrides['agnostic_nms'] = False  # NMS class-agnostic
-    model.overrides['max_det'] = 999
-    image = read_image(image)
-    results = model.predict(image)
-    render = render_result(model=model, image=image, result=results[0])
-    
-    return render
+file_urls = [
+    torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/object_detection/tree/main/samples/1.jpeg', '1.jpg'),
+    torch.hub.download_url_to_file('https://huggingface.co/spaces/foduucom/object_detection/tree/main/samples/2.JPG', '2.jpg'),
+   
 
-def gr_postprocess(image):
-    """
-    Gradio postprocess function
-    Args:
-        image: Input image
-    Returns:
-        Processed image
-    """
+]
 
-    # Convert the image to RGB
-    image = image.convert('RGB')
+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)
 
-    # Resize the image to the desired size
-    image = image.resize((640, 480))
+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"
+        )
 
-    # Convert the image to a numpy array
-    image = np.array(image)
+model = YOLO('best.pt')
+path  = [['image_0.jpg'], ['image_1.jpg']]
+video_path = [['video.mp4']]
 
-    return image
+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 = [
-# Images
-    gr.Examples(
-        [
-            'samples/1.jpeg',
-            'samples/2.JPG',
-        ],
-        inputs={'postprocess': gr_postprocess(inputs)},
-    ),
-    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"),
+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 app",
+    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 = "YOLOobjectdetection: Efficient ObjectDetection"
+inputs_video = [
+    gr.components.Video(type="filepath", label="Input Video"),
 
-description = "🔍 YoloTableExtract is a powerful space that utilizes YOLOv8s for accurate table detection and extraction. Whether tables are bordered or borderless, this space can effectively identify and extract them from images. For further assistance and support related to documentation or data-related issues, feel free to contact [email protected]. If you find this space helpful, please show your appreciation by liking it. ❤️����🏼"
-examples = [['1.jpeg', "foduucom/object_detection", 640, 0.25, 0.45], ['2.JPG', "foduucom/object_detection", 640, 0.25, 0.45]]
-demo_app = gr.Interface(
-    fn=yolov8_inference,
-    inputs=inputs,
-    outputs=outputs,
-    title=title,
-    description=description,
-    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="Pothole 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()