Add Oriented Bounding Box (OBB) support for rotated object detection

pyproject.toml

@@ -240,6 +240,7 @@ filterwarnings = [
 
 [tool.codespell]
 skip = "*.pth"
+ignore-words-list = "dota"
 
 [tool.mypy]
 python_version = "3.10"
@@ -260,6 +261,7 @@ overrides = [
   ], ignore_errors = true },
   { module = [
     "torchvision", "torchvision.*",
+    "albumentations", "albumentations.*",
     "pycocotools", "pycocotools.*",
     "timm", "timm.*",
     "einops", "einops.*",

src/rfdetr/_namespace.py

@@ -41,6 +41,7 @@
     "num_queries",
     "num_select",
     "num_windows",
+    "oriented",
     "out_feature_indexes",
     "patch_size",
     "positional_encoding_size",

src/rfdetr/config.py

@@ -515,6 +515,7 @@ class ModelConfig(BaseConfig):
     ia_bce_loss: bool = True
     cls_loss_coef: float = 1.0
     segmentation_head: bool = False
+    oriented: bool = False
     use_grouppose_keypoints: bool = False
     keypoint_cross_attn: bool = True
     inter_instance_kp_attn: bool = False
@@ -1074,9 +1075,10 @@ class TrainConfig(BaseConfig):
     keypoint_visible_loss_coef: float = 0
     keypoint_nll_loss_coef: float = 0
     keypoint_oks_sigmas: list[float] | None = None
-    dataset_file: Literal["coco", "o365", "roboflow", "yolo"] = "roboflow"
+    dataset_file: Literal["coco", "o365", "roboflow", "yolo", "dota"] = "roboflow"
     square_resize_div_64: bool = True
     dataset_dir: PathLikeStr | None
+    dota_include_difficult: bool = False
     output_dir: PathLikeStr = "output"
     multi_scale: bool = True
     expanded_scales: bool = True

src/rfdetr/datasets/__init__.py

@@ -25,6 +25,8 @@
 from rfdetr.datasets._keypoint_schema import infer_coco_keypoint_schema as infer_coco_keypoint_schema
 from rfdetr.datasets._keypoint_schema import infer_yolo_keypoint_schema as infer_yolo_keypoint_schema
 from rfdetr.datasets.coco import build_coco, build_roboflow_from_coco
+from rfdetr.datasets.dota_detection import DOTA_V1_CLASSES as DOTA_V1_CLASSES
+from rfdetr.datasets.dota_detection import DotaDetection, build_dota
 from rfdetr.datasets.o365 import build_o365
 from rfdetr.datasets.yolo import YoloDetection, build_roboflow_from_yolo
 
@@ -37,6 +39,8 @@ def get_coco_api_from_dataset(dataset: Dataset[Any]) -> Any | None:
         return dataset.coco
     if isinstance(dataset, YoloDetection):
         return dataset.coco
+    if isinstance(dataset, DotaDetection):
+        return None
     return None
 
 
@@ -95,4 +99,6 @@ def build_dataset(image_set: str, args: Any, resolution: int) -> Dataset[Any]:
         return build_roboflow(image_set, args, resolution)
     if args.dataset_file == "yolo":
         return build_roboflow_from_yolo(image_set, args, resolution)
+    if args.dataset_file == "dota":
+        return build_dota(image_set, args, resolution)
     raise ValueError(f"dataset {args.dataset_file} not supported")