Merge pull request #105 from roboflow/add-class-names

Add class names to args in checkpoint

Author: SkalskiP

rfdetr/config.py

@@ -79,3 +79,4 @@ class TrainConfig(BaseModel):
     wandb: bool = False
     project: Optional[str] = None
     run: Optional[str] = None
+    class_names: List[str] = None

rfdetr/detr.py

@@ -44,6 +44,7 @@ def train_from_config(self, config: TrainConfig, **kwargs):
         ) as f:
             anns = json.load(f)
             num_classes = len(anns["categories"])
+            class_names = [c["name"] for c in anns["categories"] if c["supercategory"] != "none"]
 
         if self.model_config.num_classes != num_classes:
             logger.warning(
@@ -52,9 +53,16 @@ def train_from_config(self, config: TrainConfig, **kwargs):
             )
             self.model.reinitialize_detection_head(num_classes)
 
+        
         train_config = config.dict()
         model_config = self.model_config.dict()
         model_config.pop("num_classes")
+        if "class_names" in model_config:
+            model_config.pop("class_names")
+        
+        if "class_names" in train_config and train_config["class_names"] is None:
+            train_config["class_names"] = class_names
+
         for k, v in train_config.items():
             if k in model_config:
                 model_config.pop(k)