merve HF staff hysts HF staff commited on
Commit
149b7f2
1 Parent(s): 9c14dad

Use gr.AnnotatedImage (#1)

Browse files

- Use gr.AnnotatedImage (d758ed068747f853c3c1918c0f5f402ad1816f79)


Co-authored-by: hysts <[email protected]>

Files changed (2) hide show
  1. app.py +18 -29
  2. requirements.txt +0 -1
app.py CHANGED
@@ -1,7 +1,5 @@
1
  import torch
2
- import cv2
3
  import gradio as gr
4
- import numpy as np
5
  from transformers import OwlViTProcessor, OwlViTForObjectDetection
6
 
7
 
@@ -26,52 +24,43 @@ def query_image(img, text_queries, score_threshold):
26
 
27
  with torch.no_grad():
28
  outputs = model(**inputs)
29
-
30
  outputs.logits = outputs.logits.cpu()
31
- outputs.pred_boxes = outputs.pred_boxes.cpu()
32
  results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes)
33
  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
34
 
35
- font = cv2.FONT_HERSHEY_SIMPLEX
36
-
37
  for box, score, label in zip(boxes, scores, labels):
38
  box = [int(i) for i in box.tolist()]
39
-
40
- if score >= score_threshold:
41
- img = cv2.rectangle(img, box[:2], box[2:], (255,0,0), 5)
42
- if box[3] + 25 > 768:
43
- y = box[3] - 10
44
- else:
45
- y = box[3] + 25
46
-
47
- img = cv2.putText(
48
- img, text_queries[label], (box[0], y), font, 1, (255,0,0), 2, cv2.LINE_AA
49
- )
50
- return img
51
 
52
 
53
  description = """
54
- Try this demo for <a href="https://huggingface.co/docs/transformers/main/en/model_doc/owlv2">OWLv2</a>,
55
- introduced in <a href="https://arxiv.org/abs/2306.09683">Scaling Open-Vocabulary Object Detection</a>.
56
- \n\n Compared to OWLVIT, OWLv2 performs better both in yield and performance (average precision).
57
- You can use OWLv2 to query images with text descriptions of any object.
58
  To use it, simply upload an image and enter comma separated text descriptions of objects you want to query the image for. You
59
- can also use the score threshold slider to set a threshold to filter out low probability predictions.
60
  \n\nOWL-ViT is trained on text templates,
61
- hence you can get better predictions by querying the image with text templates used in training the original model: e.g. *"photo of a star-spangled banner"*,
62
  *"image of a shoe"*. Refer to the <a href="https://arxiv.org/abs/2103.00020">CLIP</a> paper to see the full list of text templates used to augment the training data.
63
  \n\n<a href="https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/zeroshot_object_detection_with_owlvit.ipynb">Colab demo</a>
64
  """
65
  demo = gr.Interface(
66
- query_image,
67
- inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)],
68
- outputs="image",
69
  title="Zero-Shot Object Detection with OWLv2",
70
  description=description,
71
  examples=[
72
- ["assets/astronaut.png", "human face, rocket, star-spangled banner, nasa badge", 0.11],
73
  ["assets/coffee.png", "coffee mug, spoon, plate", 0.1],
74
  ["assets/butterflies.jpeg", "orange butterfly", 0.3],
75
  ],
76
  )
77
- demo.launch()
 
1
  import torch
 
2
  import gradio as gr
 
3
  from transformers import OwlViTProcessor, OwlViTForObjectDetection
4
 
5
 
 
24
 
25
  with torch.no_grad():
26
  outputs = model(**inputs)
27
+
28
  outputs.logits = outputs.logits.cpu()
29
+ outputs.pred_boxes = outputs.pred_boxes.cpu()
30
  results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes)
31
  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
32
 
33
+ result_labels = []
 
34
  for box, score, label in zip(boxes, scores, labels):
35
  box = [int(i) for i in box.tolist()]
36
+ if score < score_threshold:
37
+ continue
38
+ result_labels.append((box, text_queries[label.item()]))
39
+ return img, result_labels
 
 
 
 
 
 
 
 
40
 
41
 
42
  description = """
43
+ Try this demo for <a href="https://huggingface.co/docs/transformers/main/en/model_doc/owlv2">OWLv2</a>,
44
+ introduced in <a href="https://arxiv.org/abs/2306.09683">Scaling Open-Vocabulary Object Detection</a>.
45
+ \n\n Compared to OWLVIT, OWLv2 performs better both in yield and performance (average precision).
46
+ You can use OWLv2 to query images with text descriptions of any object.
47
  To use it, simply upload an image and enter comma separated text descriptions of objects you want to query the image for. You
48
+ can also use the score threshold slider to set a threshold to filter out low probability predictions.
49
  \n\nOWL-ViT is trained on text templates,
50
+ hence you can get better predictions by querying the image with text templates used in training the original model: e.g. *"photo of a star-spangled banner"*,
51
  *"image of a shoe"*. Refer to the <a href="https://arxiv.org/abs/2103.00020">CLIP</a> paper to see the full list of text templates used to augment the training data.
52
  \n\n<a href="https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/zeroshot_object_detection_with_owlvit.ipynb">Colab demo</a>
53
  """
54
  demo = gr.Interface(
55
+ query_image,
56
+ inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)],
57
+ outputs="annotatedimage",
58
  title="Zero-Shot Object Detection with OWLv2",
59
  description=description,
60
  examples=[
61
+ ["assets/astronaut.png", "human face, rocket, star-spangled banner, nasa badge", 0.11],
62
  ["assets/coffee.png", "coffee mug, spoon, plate", 0.1],
63
  ["assets/butterflies.jpeg", "orange butterfly", 0.3],
64
  ],
65
  )
66
+ demo.launch()
requirements.txt CHANGED
@@ -2,5 +2,4 @@ numpy>=1.18.5
2
  torch>=1.7.0
3
  torchvision>=0.8.1
4
  git+https://github.com/nielsrogge/transformers@add_owlv2
5
- opencv-python
6
  scipy
 
2
  torch>=1.7.0
3
  torchvision>=0.8.1
4
  git+https://github.com/nielsrogge/transformers@add_owlv2
 
5
  scipy