Optimize classification models

focoos/data/catalog/catalog.py

@@ -85,6 +85,19 @@ class CatalogDataset:
             filter_empty=False,
         ),
     ),
+    CatalogDataset(
+        name="coco_2017_cls",
+        task=Task.CLASSIFICATION,
+        train_split=CatalogSplit(
+            image_root="coco/train2017",
+            json_file="coco/annotations/instances_train2017.json",
+        ),
+        val_split=CatalogSplit(
+            image_root="coco/val2017",
+            json_file="coco/annotations/instances_val2017.json",
+            filter_empty=False,
+        ),
+    ),
     CatalogDataset(
         name="coco_2017_instance",
         task=Task.INSTANCE_SEGMENTATION,
@@ -131,6 +144,7 @@ def _load_dataset_split(
     split_name: str,
     split: CatalogSplit,
     task: Task,
+    split_type: DatasetSplitType,
     root=DATASETS_DIR,
 ) -> DictDataset:
     """
@@ -152,7 +166,7 @@ def get_path(root, path):
         task=task,
     )
 
-    if task in [Task.DETECTION, Task.INSTANCE_SEGMENTATION, Task.KEYPOINT]:
+    if task in [Task.DETECTION, Task.INSTANCE_SEGMENTATION, Task.KEYPOINT, Task.CLASSIFICATION]:
         dataset_dict = load_coco_json(json_file_path, image_root_path, metadata, task=task)
         if split.filter_empty:
             dataset_dict = filter_images_with_only_crowd_annotations(dataset_dicts=dataset_dict)
@@ -171,7 +185,7 @@ def get_path(root, path):
         raise ValueError(f"Unknown task {task}")
 
     metadata.count = len(dataset_dict)
-    return DictDataset(dataset_dict, task=task, metadata=metadata)
+    return DictDataset(dataset_dict, task=task, metadata=metadata, split_type=split_type)
 
 
 def get_dataset_split(name: str, split: DatasetSplitType, datasets_root=DATASETS_DIR) -> DictDataset:
@@ -192,4 +206,4 @@ def get_dataset_split(name: str, split: DatasetSplitType, datasets_root=DATASETS
     else:
         raise ValueError(f"Unknown split {split}")
 
-    return _load_dataset_split(split_name, entry, ds.task, datasets_root)
+    return _load_dataset_split(split_name=split_name, split=entry, task=ds.task, split_type=split, root=datasets_root)

focoos/infer/runtimes/torchscript.py

@@ -49,7 +49,9 @@ def __init__(
 
         if self.opts.warmup_iter > 0:
             size = (
-                self.model_info.im_size if self.model_info.task == Task.DETECTION and self.model_info.im_size else 640
+                self.model_info.im_size
+                if self.model_info.task in [Task.DETECTION, Task.CLASSIFICATION] and self.model_info.im_size
+                else 640
             )
             logger.info(f"⏱️ Warming up model {self.model_info.name} on {self.device}, size: {size}x{size}..")
             with torch.no_grad():