SkalskiP commited on
Commit
30841fb
1 Parent(s): 7500bfd

Examples panel added

Browse files
Files changed (2) hide show
  1. app.py +18 -3
  2. 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.003,
36
  iou_threshold: float = 0.5,
37
  with_segmentation: bool = True,
38
- with_confidence: bool = True
 
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).with_nms(iou_threshold)
 
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.12
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