💬 [Update] Bbox drawing tools, change label text

yolo/config/config.py

@@ -140,6 +140,7 @@ class Config:
     cpu_num: int
 
     class_num: int
+    class_list: List[str]
     image_size: List[int]
 
     out_path: str

yolo/config/general.yaml

@@ -2,6 +2,7 @@ device: 0
 cpu_num: 16
 
 class_num: 80
+class_list: ['Person', 'Bicycle', 'Car', 'Motorcycle', 'Airplane', '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', 'Couch', 'Potted plant', 'Bed', 'Dining table', 'Toilet', 'Tv', 'Laptop', 'Mouse', 'Remote', 'Keyboard', 'Cell phone', 'Microwave', 'Oven', 'Toaster', 'Sink', 'Refrigerator', 'Book', 'Clock', 'Vase', 'Scissors', 'Teddy bear', 'Hair drier', 'Toothbrush']
 image_size: [640, 640]
 
 out_path: runs

yolo/tools/data_augmentation.py

@@ -9,7 +9,7 @@ class AugmentationComposer:
 
     def __init__(self, transforms, image_size: int = 640):
         self.transforms = transforms
-        self.image_size = image_size
+        self.image_size = image_size[0]
         self.pad_resize = PadAndResize(self.image_size)
 
         for transform in self.transforms:

yolo/tools/solver.py

@@ -106,6 +106,7 @@ class ModelTester:
 
         self.anchor2box = AnchorBoxConverter(cfg.model, cfg.image_size, device)
         self.nms = cfg.task.nms
+        self.idx2label = cfg.class_list
         self.save_path = save_path
 
     def solve(self, dataloader: StreamDataLoader):
@@ -119,7 +120,12 @@ class ModelTester:
                 predict, _ = self.anchor2box(raw_output[0][3:], with_logits=True)
                 nms_out = bbox_nms(predict, self.nms)
                 draw_bboxes(
-                    images[0], nms_out[0], scaled_bbox=False, save_path=self.save_path, save_name=f"frame{idx:03d}.png"
+                    images[0],
+                    nms_out[0],
+                    scaled_bbox=False,
+                    save_path=self.save_path,
+                    save_name=f"frame{idx:03d}.png",
+                    idx2label=self.idx2label,
                 )
         except (KeyboardInterrupt, Exception) as e:
             dataloader.stop_event.set()

yolo/utils/bounding_box_utils.py

@@ -307,7 +307,7 @@ def bbox_nms(predicts: Tensor, nms_cfg: NMSConfig):
         instance_idx = nms_idx[idx == batch_idx[nms_idx]]
 
         predict_nms = torch.cat(
-            [valid_cls[instance_idx][:, None], valid_con[instance_idx][:, None], valid_box[instance_idx]], dim=-1
+            [valid_cls[instance_idx][:, None], valid_box[instance_idx], valid_con[instance_idx][:, None]], dim=-1
         )
 
         predicts_nms.append(predict_nms)
@@ -322,7 +322,7 @@ def calculate_map(predictions, ground_truths, iou_thresholds):
     ground_truths = ground_truths[:n_gts]
     aps = []
 
-    ious = calculate_iou(predictions[:, 2:], ground_truths[:, 1:])  # [n_preds, n_gts]
+    ious = calculate_iou(predictions[:, 1:-1], ground_truths[:, 1:])  # [n_preds, n_gts]
 
     for threshold in iou_thresholds:
         tp = torch.zeros(n_preds, device=device)