Update app.py

app.py

@@ -4,10 +4,9 @@ import torch
 import supervision as sv
 import json
 
-def anylize(img):
-    id2label = {0: 'dangerous-items', 1: 'Gun', 2: 'Knife', 3: 'Pliers', 4: 'Scissors', 5: 'Wrench'}
-    image = img
+id2label = {0: 'dangerous-items', 1: 'Gun', 2: 'Knife', 3: 'Pliers', 4: 'Scissors', 5: 'Wrench'}
 
+def anylize(image):
     with torch.no_grad():
 
         inputs = image_processor(images=image, return_tensors='pt')
@@ -22,7 +21,23 @@ def anylize(img):
 
     # annotate
     detections = sv.Detections.from_transformers(transformers_results=results).with_nms(threshold=0.5)
-    labels = [str([list(xyxy), confidence, id2label[class_id]]) for xyxy, _, confidence, class_id, _ in detections]
-    json_list = json.dumps(str(labels[0]))
 
-gr.Interface.load("models/Guy2/AirportSec-100epoch").launch()
+    out = {}
+    
+    for idx, detection in enumerate(detections):
+        cls = id2label[detection.class_id]
+        confidence = detection.confidence
+        box = detection.xyxy
+
+        out[str(idx)] = {
+            "box":list(box),
+            "cls":cls,
+            "conf":confidence
+        }
+        
+    #labels = [str([list(xyxy), confidence, id2label[class_id]]) for xyxy, _, confidence, class_id, _ in detections]
+    #json_list = json.dumps(str(labels[0]))
+
+    return out
+
+gr.Interface(fn = anylize, inputs="image", outputs=gr.JSON()).launch()