codingwithlewis commited on
Commit
7d0b9c0
1 Parent(s): e413988

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -6
app.py CHANGED
@@ -2,6 +2,7 @@ import torch
2
  import gradio as gr
3
  from transformers import Owlv2Processor, Owlv2ForObjectDetection
4
  import spaces
 
5
 
6
  # Use GPU if available
7
  if torch.cuda.is_available():
@@ -30,12 +31,18 @@ def query_image(img, text_queries, score_threshold):
30
  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
31
 
32
  result_labels = []
 
33
  for box, score, label in zip(boxes, scores, labels):
34
  box = [int(i) for i in box.tolist()]
35
  if score < score_threshold:
36
  continue
 
 
 
 
37
  result_labels.append((box, text_queries[label.item()]))
38
- return img, result_labels, boxes
 
39
 
40
 
41
  description = """
@@ -56,10 +63,5 @@ demo = gr.Interface(
56
  outputs=["annotatedimage", "json"],
57
  title="Zero-Shot Object Detection with OWLv2",
58
  description=description,
59
- examples=[
60
- ["assets/astronaut.png", "human face, rocket, star-spangled banner, nasa badge", 0.11],
61
- ["assets/coffee.png", "coffee mug, spoon, plate", 0.1],
62
- ["assets/butterflies.jpeg", "orange butterfly", 0.3],
63
- ],
64
  )
65
  demo.launch()
 
2
  import gradio as gr
3
  from transformers import Owlv2Processor, Owlv2ForObjectDetection
4
  import spaces
5
+ import json
6
 
7
  # Use GPU if available
8
  if torch.cuda.is_available():
 
31
  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
32
 
33
  result_labels = []
34
+ boxes_coords = []
35
  for box, score, label in zip(boxes, scores, labels):
36
  box = [int(i) for i in box.tolist()]
37
  if score < score_threshold:
38
  continue
39
+ boxes_coords.append({
40
+ "object": text_queries[label.item()],
41
+ "pos": box
42
+ })
43
  result_labels.append((box, text_queries[label.item()]))
44
+ print(boxes_coords)
45
+ return [img, result_labels], boxes_coords
46
 
47
 
48
  description = """
 
63
  outputs=["annotatedimage", "json"],
64
  title="Zero-Shot Object Detection with OWLv2",
65
  description=description,
 
 
 
 
 
66
  )
67
  demo.launch()