v04 add color

app.py

@@ -160,23 +160,25 @@ def export_json(results, img_size):
 
 
 # frame conversion
-def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt):
+def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt, obj_cls_index, color_list):
 
     img_pil = ImageDraw.Draw(img)
 
-    img_pil.rectangle(xyxy, fill=None, outline="green")  # bounding box
+    img_pil.rectangle(xyxy, fill=None, outline=color_list[obj_cls_index])  # bounding box
 
     if "label" in opt:
         text_w, text_h = textFont.getsize(countdown_msg)  # Label size
+        
         img_pil.rectangle(
             (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h),
-            fill="green",
-            outline="green",
+            fill=color_list[obj_cls_index],
+            outline=color_list[obj_cls_index],
         )  # label background
+        
         img_pil.multiline_text(
             (xyxy[0], xyxy[1]),
             countdown_msg,
-            fill=(205, 250, 255),
+            fill=(255, 255, 255),
             font=textFont,
             align="center",
         )
@@ -184,6 +186,16 @@ def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt):
     return img
 
 
+def color_set(cls_num):
+    color_list = []
+    for i in range(cls_num):
+        color = tuple(np.random.choice(range(256), size=3))
+        # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
+        color_list.append(color)
+    
+    return color_list
+
+
 # YOLOv5 image detection function
 def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_cls, opt):
 
@@ -210,6 +222,8 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
     model.max_det = int(max_num)  # Maximum number of detection frames
     model.classes = model_cls  # model classes
 
+    color_list = color_set(len(model_cls_name_cp)) # 设置颜色
+
     img_size = img.size  # frame size
 
     results = model(img, size=infer_size)  # detection
@@ -260,6 +274,8 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
                 [x0, y0, x1, y1],
                 FONTSIZE,
                 opt,
+                obj_cls_index,
+                color_list,
             )
 
             # ----------add object size----------
@@ -332,6 +348,8 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
     model.max_det = int(max_num)  # Maximum number of detection frames
     model.classes = model_cls  # model classes
 
+    color_list = color_set(len(model_cls_name_cp)) # 设置颜色
+
     # ----------------Load fonts----------------
     yaml_index = cls_name.index(".yaml")
     cls_name_lang = cls_name[yaml_index - 2:yaml_index]
@@ -393,6 +411,8 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
                     [x0, y0, x1, y1],
                     FONTSIZE,
                     opt,
+                    obj_cls_index,
+                    color_list,
                 )
 
                 frame = cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR)