added labels

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 import numpy as np
+import yaml
 
 from yolov5.models.common import DetectMultiBackend
 from yolov5.utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
@@ -23,6 +24,9 @@ hide_labels = False
 hide_conf = True
 line_thickness = 1
 
+with open('yolov5/joint_all_multi.yaml', 'r') as f:
+    LABELS = yaml.safe_load(f)['names']
+
 def joint_detection(img0):
     global imgsz
     img = letterbox(img0, 640, stride=stride, auto=pt)[0]
@@ -66,9 +70,23 @@ def joint_detection(img0):
                 label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                 annotator.box_label(xyxy, label, color=colors(c, True))
 
+        # save as text
+        # Write results
+        content = {}
+        for *xyxy, conf, cls in reversed(det):
+            xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
+            x, y, width, height = xywh
+            current_label = LABELS[int(cls.item())]
+
+            if content.get(current_label, None) is None:
+                content[current_label] = []
+
+            current_dict = {'x': x, 'y': y, 'width': width, 'height': height}
+            content[current_label].append(current_dict)  # label format
+
         # Stream results
         im0 = annotator.result()
-        return im0
+        return im0, content
         # if view_img:
         #     cv2.imshow(str(p), im0)
         #     cv2.waitKey(1)  # 1 millisecond
@@ -94,5 +112,5 @@ def joint_detection(img0):
 
     # Print time (inference-only)
 
-iface = gr.Interface(fn=joint_detection, inputs="image", outputs="image")
+iface = gr.Interface(fn=joint_detection, inputs="image", outputs=["image", "json"])
 iface.launch()

yolov5/joint_all_multi.yaml

@@ -0,0 +1,8 @@
+path: ../datasets/joint_all_multi  # dataset root dir
+train: train.txt  # train images (relative to 'path') 118287 images
+val: eval.txt  # train images (relative to 'path') 5000 images
+#test: test-dev2017.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
+
+# Classes
+nc: 5  # number of classes
+names: ['PIP', 'MCP', 'Wrist', 'Radius', 'Ulna']