adirik commited on
Commit
bc5dfe0
1 Parent(s): dff85a2

fix styling

Browse files
Files changed (3) hide show
  1. .DS_Store +0 -0
  2. app.py +15 -9
  3. requirements.txt +2 -1
.DS_Store CHANGED
Binary files a/.DS_Store and b/.DS_Store differ
 
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import torch
 
2
  import gradio as gr
3
  import numpy as np
4
  from PIL import Image, ImageDraw, ImageFont
@@ -19,19 +20,24 @@ def query_image(img, text_queries):
19
  results = processor.post_process(outputs=outputs, target_sizes=target_sizes)
20
  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
21
 
22
- draw = ImageDraw.Draw(img)
23
- font = ImageFont.truetype("assets/Helvetica.ttf", size=22)
 
 
24
 
25
- score_threshold = 0.1
26
  for box, score, label in zip(boxes, scores, labels):
27
  box = [int(i) for i in box.tolist()]
28
 
29
  if score >= score_threshold:
30
- draw.rectangle(box, outline="red", width=4)
31
- text_loc =[box[0]+5, box[3]+10]
32
- draw.text(text_loc, text_queries[label], fill="red", font=font, stroke_width=1)
33
-
34
- img = np.array(img)
 
 
 
 
35
  return img
36
 
37
 
@@ -45,7 +51,7 @@ To use it, simply upload an image and enter comma separated text descriptions of
45
  """
46
  demo = gr.Interface(
47
  query_image,
48
- inputs=[gr.Image(shape=(768, 768), type="pil"), "text"],
49
  outputs="image",
50
  title="Zero-Shot Object Detection with OWL-ViT",
51
  description=description,
 
1
  import torch
2
+ import cv2
3
  import gradio as gr
4
  import numpy as np
5
  from PIL import Image, ImageDraw, ImageFont
 
20
  results = processor.post_process(outputs=outputs, target_sizes=target_sizes)
21
  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
22
 
23
+ img = cv2.resize(img, (768, 768), interpolation = cv2.INTER_AREA)
24
+ score_threshold = 0.11
25
+
26
+ font = cv2.FONT_HERSHEY_SIMPLEX
27
 
 
28
  for box, score, label in zip(boxes, scores, labels):
29
  box = [int(i) for i in box.tolist()]
30
 
31
  if score >= score_threshold:
32
+ img = cv2.rectangle(img, box[:2], box[2:], (255,0,0), 5)
33
+ if box[3] + 25 > 768:
34
+ y = box[3] - 10
35
+ else:
36
+ y = box[3] + 25
37
+
38
+ img = cv2.putText(
39
+ img, text_queries[label], (box[0], y), font, 1, (255,0,0), 2, cv2.LINE_AA
40
+ )
41
  return img
42
 
43
 
 
51
  """
52
  demo = gr.Interface(
53
  query_image,
54
+ inputs=[gr.Image(shape=(768, 768)), "text"],
55
  outputs="image",
56
  title="Zero-Shot Object Detection with OWL-ViT",
57
  description=description,
requirements.txt CHANGED
@@ -4,4 +4,5 @@ numpy>=1.18.5
4
  torch>=1.7.0
5
  torchvision>=0.8.1
6
  Pillow
7
- transformers
 
 
4
  torch>=1.7.0
5
  torchvision>=0.8.1
6
  Pillow
7
+ transformers
8
+ opencv-python