Spaces:
Runtime error
Runtime error
Examples panel added
Browse files- app.py +18 -3
- requirements.txt +1 -1
app.py
CHANGED
|
@@ -1,5 +1,6 @@
|
|
| 1 |
from typing import List
|
| 2 |
|
|
|
|
| 3 |
import torch
|
| 4 |
import gradio as gr
|
| 5 |
import numpy as np
|
|
@@ -16,6 +17,10 @@ This is a demo of zero-shot instance segmentation using [YOLO-World](https://git
|
|
| 16 |
Powered by Roboflow [Inference](https://github.com/roboflow/inference) and [Supervision](https://github.com/roboflow/supervision).
|
| 17 |
"""
|
| 18 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 20 |
EFFICIENT_SAM_MODEL = load(device=DEVICE)
|
| 21 |
YOLO_WORLD_MODEL = YOLOWorld(model_id="yolo_world/l")
|
|
@@ -32,15 +37,17 @@ def process_categories(categories: str) -> List[str]:
|
|
| 32 |
def process_image(
|
| 33 |
input_image: np.ndarray,
|
| 34 |
categories: str,
|
| 35 |
-
confidence_threshold: float = 0.
|
| 36 |
iou_threshold: float = 0.5,
|
| 37 |
with_segmentation: bool = True,
|
| 38 |
-
with_confidence: bool =
|
|
|
|
| 39 |
) -> np.ndarray:
|
| 40 |
categories = process_categories(categories)
|
| 41 |
YOLO_WORLD_MODEL.set_classes(categories)
|
| 42 |
results = YOLO_WORLD_MODEL.infer(input_image, confidence=confidence_threshold)
|
| 43 |
-
detections = sv.Detections.from_inference(results)
|
|
|
|
| 44 |
if with_segmentation:
|
| 45 |
masks = []
|
| 46 |
for [x_min, y_min, x_max, y_max] in detections.xyxy:
|
|
@@ -55,9 +62,11 @@ def process_image(
|
|
| 55 |
zip(detections.class_id, detections.confidence)
|
| 56 |
]
|
| 57 |
output_image = input_image.copy()
|
|
|
|
| 58 |
output_image = MASK_ANNOTATOR.annotate(output_image, detections)
|
| 59 |
output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
|
| 60 |
output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels)
|
|
|
|
| 61 |
return output_image
|
| 62 |
|
| 63 |
|
|
@@ -79,6 +88,12 @@ with gr.Blocks() as demo:
|
|
| 79 |
scale=5
|
| 80 |
)
|
| 81 |
submit_button_component = gr.Button('Submit', scale=1)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 82 |
|
| 83 |
submit_button_component.click(
|
| 84 |
fn=process_image,
|
|
|
|
| 1 |
from typing import List
|
| 2 |
|
| 3 |
+
import cv2
|
| 4 |
import torch
|
| 5 |
import gradio as gr
|
| 6 |
import numpy as np
|
|
|
|
| 17 |
Powered by Roboflow [Inference](https://github.com/roboflow/inference) and [Supervision](https://github.com/roboflow/supervision).
|
| 18 |
"""
|
| 19 |
|
| 20 |
+
EXAMPLES = [
|
| 21 |
+
['https://media.roboflow.com/dog.jpeg', 'dog, eye, nose, tongue, car', 0.005, 0.1, True, False, False],
|
| 22 |
+
]
|
| 23 |
+
|
| 24 |
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 25 |
EFFICIENT_SAM_MODEL = load(device=DEVICE)
|
| 26 |
YOLO_WORLD_MODEL = YOLOWorld(model_id="yolo_world/l")
|
|
|
|
| 37 |
def process_image(
|
| 38 |
input_image: np.ndarray,
|
| 39 |
categories: str,
|
| 40 |
+
confidence_threshold: float = 0.005,
|
| 41 |
iou_threshold: float = 0.5,
|
| 42 |
with_segmentation: bool = True,
|
| 43 |
+
with_confidence: bool = False,
|
| 44 |
+
with_class_agnostic_nms: bool = False,
|
| 45 |
) -> np.ndarray:
|
| 46 |
categories = process_categories(categories)
|
| 47 |
YOLO_WORLD_MODEL.set_classes(categories)
|
| 48 |
results = YOLO_WORLD_MODEL.infer(input_image, confidence=confidence_threshold)
|
| 49 |
+
detections = sv.Detections.from_inference(results)
|
| 50 |
+
detections = detections.with_nms(class_agnostic=with_class_agnostic_nms, threshold=iou_threshold)
|
| 51 |
if with_segmentation:
|
| 52 |
masks = []
|
| 53 |
for [x_min, y_min, x_max, y_max] in detections.xyxy:
|
|
|
|
| 62 |
zip(detections.class_id, detections.confidence)
|
| 63 |
]
|
| 64 |
output_image = input_image.copy()
|
| 65 |
+
output_image = cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR)
|
| 66 |
output_image = MASK_ANNOTATOR.annotate(output_image, detections)
|
| 67 |
output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
|
| 68 |
output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels)
|
| 69 |
+
output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)
|
| 70 |
return output_image
|
| 71 |
|
| 72 |
|
|
|
|
| 88 |
scale=5
|
| 89 |
)
|
| 90 |
submit_button_component = gr.Button('Submit', scale=1)
|
| 91 |
+
gr.Examples(
|
| 92 |
+
fn=process_image,
|
| 93 |
+
examples=EXAMPLES,
|
| 94 |
+
inputs=[input_image_component, categories_text_component],
|
| 95 |
+
outputs=output_image_component
|
| 96 |
+
)
|
| 97 |
|
| 98 |
submit_button_component.click(
|
| 99 |
fn=process_image,
|
requirements.txt
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
-
inference-gpu[yolo-world]==0.9.
|
| 2 |
supervision==0.19.0rc3
|
| 3 |
gradio==4.19.0
|
|
|
|
| 1 |
+
inference-gpu[yolo-world]==0.9.13
|
| 2 |
supervision==0.19.0rc3
|
| 3 |
gradio==4.19.0
|