rdezwart commited on
Commit
da68fcd
·
1 Parent(s): 927201f

Try creating cropped images from bounding boxes

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -57,12 +57,16 @@ def detect_objects(img: Image.Image):
57
  target_sizes = torch.tensor([tuple(reversed(img.size))])
58
  results = yolos_processor.post_process_object_detection(outputs, threshold=0.7, target_sizes=target_sizes)[0]
59
 
 
60
  for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
61
  box = [round(i, 2) for i in box.tolist()]
62
  print(
63
  f"Detected {yolos_model.config.id2label[label.item()]} with confidence "
64
  f"{round(score.item(), 3)} at location {box}"
65
  )
 
 
 
66
 
67
 
68
  if __name__ == "__main__":
@@ -90,6 +94,6 @@ if __name__ == "__main__":
90
  moon_output = gr.TextArea(label="Output")
91
 
92
  moon_submit.click(answer_question, [moon_img, moon_prompt], moon_output)
93
- yolos_button.click(detect_objects, [yolos_input])
94
 
95
  app.queue().launch()
 
57
  target_sizes = torch.tensor([tuple(reversed(img.size))])
58
  results = yolos_processor.post_process_object_detection(outputs, threshold=0.7, target_sizes=target_sizes)[0]
59
 
60
+ box_images = []
61
  for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
62
  box = [round(i, 2) for i in box.tolist()]
63
  print(
64
  f"Detected {yolos_model.config.id2label[label.item()]} with confidence "
65
  f"{round(score.item(), 3)} at location {box}"
66
  )
67
+ box_images.append(img[box[1]:box[3], box[0]:box[2]])
68
+
69
+ return box_images[0]
70
 
71
 
72
  if __name__ == "__main__":
 
94
  moon_output = gr.TextArea(label="Output")
95
 
96
  moon_submit.click(answer_question, [moon_img, moon_prompt], moon_output)
97
+ yolos_button.click(detect_objects, [yolos_input], yolos_output)
98
 
99
  app.queue().launch()