Guy2 commited on
Commit
77c3c13
·
1 Parent(s): 6503948
Files changed (1) hide show
  1. app.py +4 -6
app.py CHANGED
@@ -2,19 +2,17 @@ import gradio as gr
2
  from transformers import DetrImageProcessor, DetrForObjectDetection
3
  import torch
4
  import supervision as sv
 
5
 
6
  def anylize(img):
7
- id2label = {k: v['name'] for k,v in categories.items()}
8
- box_annotator = sv.BoxAnnotator()
9
  image = img
10
 
11
  with torch.no_grad():
12
 
13
- # load image and predict
14
  inputs = image_processor(images=image, return_tensors='pt')
15
  outputs = model(**inputs)
16
 
17
- # post-process
18
  target_sizes = torch.tensor([image.shape[:2]])
19
  results = image_processor.post_process_object_detection(
20
  outputs=outputs,
@@ -24,7 +22,7 @@ def anylize(img):
24
 
25
  # annotate
26
  detections = sv.Detections.from_transformers(transformers_results=results).with_nms(threshold=0.5)
27
- labels = [f"{id2label[class_id]} {confidence:.2f}" for _, confidence, class_id, _ in detections]
28
- frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)
29
 
30
  gr.Interface.load("models/Guy2/AirportSec-100epoch").launch()
 
2
  from transformers import DetrImageProcessor, DetrForObjectDetection
3
  import torch
4
  import supervision as sv
5
+ import json
6
 
7
  def anylize(img):
8
+ id2label = {0: 'dangerous-items', 1: 'Gun', 2: 'Knife', 3: 'Pliers', 4: 'Scissors', 5: 'Wrench'}
 
9
  image = img
10
 
11
  with torch.no_grad():
12
 
 
13
  inputs = image_processor(images=image, return_tensors='pt')
14
  outputs = model(**inputs)
15
 
 
16
  target_sizes = torch.tensor([image.shape[:2]])
17
  results = image_processor.post_process_object_detection(
18
  outputs=outputs,
 
22
 
23
  # annotate
24
  detections = sv.Detections.from_transformers(transformers_results=results).with_nms(threshold=0.5)
25
+ labels = [str([list(xyxy), confidence, id2label[class_id]]) for xyxy, _, confidence, class_id, _ in detections]
26
+ json_list = json.dumps(str(labels[0]))
27
 
28
  gr.Interface.load("models/Guy2/AirportSec-100epoch").launch()