🚚 [Update] the class_num&class_list to dataset

yolo/config/config.py

@@ -45,6 +45,8 @@ class DownloadOptions:
 @dataclass
 class DatasetConfig:
     path: str
+    class_num: int
+    class_list: List[str]
     auto_download: Optional[DownloadOptions]
 
 
@@ -142,8 +144,6 @@ class Config:
     device: Union[str, int, List[int]]
     cpu_num: int
 
-    class_num: int
-    class_list: List[str]
     class_idx_id: List[int]
     image_size: List[int]
 

yolo/config/dataset/coco.yaml

@@ -2,6 +2,9 @@ path: data/coco
 train: train2017
 validation: val2017
 
+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']
+
 auto_download:
   images:
     base_url: http://images.cocodataset.org/zips/

yolo/config/dataset/dev.yaml

@@ -2,4 +2,7 @@ path: data/dev
 train: train
 validation: val
 
+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']
+
 auto_download:

yolo/config/dataset/mock.yaml

@@ -2,6 +2,9 @@ path: tests/data
 train: train
 validation: val
 
+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']
+
 auto_download:
   images:
     base_url: https://github.com/WongKinYiu/yolov9mit/releases/download/v1.0-alpha/

yolo/config/general.yaml

@@ -1,8 +1,6 @@
 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/lazy.py

@@ -24,7 +24,7 @@ def main(cfg: Config):
     if getattr(cfg.task, "fast_inference", False):
         model = FastModelLoader(cfg).load_model(device)
     else:
-        model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
+        model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
         model = model.to(device)
 
     converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)

yolo/tools/loss_functions.py

@@ -109,7 +109,7 @@ class YOLOLoss:
 class DualLoss:
     def __init__(self, cfg: Config, vec2box) -> None:
         loss_cfg = cfg.task.loss
-        self.loss = YOLOLoss(loss_cfg, vec2box, class_num=cfg.class_num, reg_max=cfg.model.anchor.reg_max)
+        self.loss = YOLOLoss(loss_cfg, vec2box, class_num=cfg.dataset.class_num, reg_max=cfg.model.anchor.reg_max)
 
         self.aux_rate = loss_cfg.aux
 

yolo/tools/solver.py

@@ -162,7 +162,7 @@ class ModelTester:
         self.save_path = progress.save_path / "images"
         os.makedirs(self.save_path, exist_ok=True)
         self.save_predict = getattr(cfg.task, "save_predict", None)
-        self.idx2label = cfg.class_list
+        self.idx2label = cfg.dataset.class_list
 
     def solve(self, dataloader: StreamDataLoader):
         logger.info("👀 Start Inference!")

yolo/utils/deploy_utils.py

@@ -12,6 +12,8 @@ class FastModelLoader:
     def __init__(self, cfg: Config):
         self.cfg = cfg
         self.compiler = cfg.task.fast_inference
+        self.class_num = cfg.dataset.class_num
+
         self._validate_compiler()
         if cfg.weight == True:
             cfg.weight = Path("weights") / f"{cfg.model.name}.pt"
@@ -32,7 +34,7 @@ class FastModelLoader:
             return self._load_trt_model().to(device)
         elif self.compiler == "deploy":
             self.cfg.model.model.auxiliary = {}
-        return create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).to(device)
+        return create_model(self.cfg.model, class_num=self.class_num, weight_path=self.cfg.weight).to(device)
 
     def _load_onnx_model(self, device):
         from onnxruntime import InferenceSession
@@ -67,7 +69,7 @@ class FastModelLoader:
         from onnxruntime import InferenceSession
         from torch.onnx import export
 
-        model = create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).eval()
+        model = create_model(self.cfg.model, class_num=self.class_num, weight_path=self.cfg.weight).eval()
         dummy_input = torch.ones((1, 3, *self.cfg.image_size))
         export(
             model,
@@ -95,7 +97,7 @@ class FastModelLoader:
     def _create_trt_model(self):
         from torch2trt import torch2trt
 
-        model = create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).eval()
+        model = create_model(self.cfg.model, class_num=self.class_num, weight_path=self.cfg.weight).eval()
         dummy_input = torch.ones((1, 3, *self.cfg.image_size)).cuda()
         logger.info(f"♻️ Creating TensorRT model")
         model_trt = torch2trt(model.cuda(), [dummy_input])