Guy2 commited on
Commit
6503948
·
1 Parent(s): d9f0b09
Files changed (1) hide show
  1. app.py +27 -0
app.py CHANGED
@@ -1,3 +1,30 @@
1
  import gradio as gr
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  gr.Interface.load("models/Guy2/AirportSec-100epoch").launch()
 
1
  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,
21
+ threshold=0.8,
22
+ target_sizes=target_sizes
23
+ )[0]
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()