Add class label to returned items

tt-metal-yolov4/server/coco.names

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorbike
+aeroplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+sofa
+pottedplant
+bed
+diningtable
+toilet
+tvmonitor
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

tt-metal-yolov4/server/fast_api_yolov4.py

@@ -43,6 +43,22 @@ def get_dispatch_core_type():
 
 @app.on_event("startup")
 async def startup():
+    global class_names
+
+    def load_class_names(namesfile):
+        class_names = []
+        with open(namesfile, "r") as fp:
+            lines = fp.readlines()
+            for line in lines:
+                line = line.rstrip()
+                class_names.append(line)
+        return class_names
+
+    namesfile = "coco.names"
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    file_path = os.path.join(script_dir, namesfile)
+    class_names = load_class_names(file_path)
+
     global model
     if ("WH_ARCH_YAML" in os.environ) and os.environ[
         "WH_ARCH_YAML"
@@ -84,6 +100,7 @@ def process_output(output):
     for item in output:
         cnt = cnt + 1
         output_i = [element.item() for element in item]
+        output_i.append(class_names[output_i[-1]])  # map class id to class name
         outs.append(output_i)
     return outs
 
@@ -134,7 +151,6 @@ def post_processing(img, conf_thresh, nms_thresh, output):
                             ll_box_array[k, 2],
                             ll_box_array[k, 3],
                             ll_max_conf[k],
-                            ll_max_conf[k],
                             ll_max_id[k],
                         ]
                     )

tt-metal-yolov4/yolov4.src.Dockerfile

@@ -91,7 +91,7 @@ ENV PYTHONPATH=${PYTHONPATH}:${APP_DIR}
 COPY --chown=user:user "/server" "${APP_DIR}/server"
 COPY --chown=user:user "/requirements.txt" "${APP_DIR}/requirements.txt"
 RUN /bin/bash -c "source ${PYTHON_ENV_DIR}/bin/activate \
-&& pip install --default-timeout=240 --no-cache-dir -r requirements.txt"
+    && pip install --default-timeout=240 --no-cache-dir -r requirements.txt"
 
 # spinup inference server
 WORKDIR "${TT_METAL_HOME}"