From ac8839a7a0361ccd3c6b4233bb0128465997c79c Mon Sep 17 00:00:00 2001
From: rentainhe <596106517@qq.com>
Date: Tue, 22 Aug 2023 17:21:24 +0800
Subject: [PATCH] add lvis dataset

---
 configs/common/data/lvis_detr.py | 74 ++++++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)
 create mode 100644 configs/common/data/lvis_detr.py

diff --git a/configs/common/data/lvis_detr.py b/configs/common/data/lvis_detr.py
new file mode 100644
index 00000000..b6a63c07
--- /dev/null
+++ b/configs/common/data/lvis_detr.py
@@ -0,0 +1,74 @@
+from omegaconf import OmegaConf
+
+import detectron2.data.transforms as T
+from detectron2.config import LazyCall as L
+from detectron2.data import (
+    build_detection_test_loader,
+    build_detection_train_loader,
+    get_detection_dataset_dicts,
+)
+from detectron2.evaluation import LVISEvaluator
+
+from detrex.data import DetrDatasetMapper
+
+dataloader = OmegaConf.create()
+
+dataloader.train = L(build_detection_train_loader)(
+    dataset=L(get_detection_dataset_dicts)(names="lvis_v1_train"),
+    mapper=L(DetrDatasetMapper)(
+        augmentation=[
+            L(T.RandomFlip)(),
+            L(T.ResizeShortestEdge)(
+                short_edge_length=(480, 512, 544, 576, 608, 640,
+                                   672, 704, 736, 768, 800),
+                max_size=1333,
+                sample_style="choice",
+            ),
+        ],
+        augmentation_with_crop=[
+            L(T.RandomFlip)(),
+            L(T.ResizeShortestEdge)(
+                short_edge_length=(400, 500, 600),
+                sample_style="choice",
+            ),
+            L(T.RandomCrop)(
+                crop_type="absolute_range",
+                crop_size=(384, 600),
+            ),
+            L(T.ResizeShortestEdge)(
+                short_edge_length=(480, 512, 544, 576, 608, 640,
+                                   672, 704, 736, 768, 800),
+                max_size=1333,
+                sample_style="choice",
+            ),
+        ],
+        is_train=True,
+        mask_on=False,
+        img_format="RGB",
+    ),
+    total_batch_size=16,
+    num_workers=4,
+)
+
+dataloader.test = L(build_detection_test_loader)(
+    dataset=L(get_detection_dataset_dicts)(names="lvis_v1_val",
+                                           filter_empty=False),
+    mapper=L(DetrDatasetMapper)(
+        augmentation=[
+            L(T.ResizeShortestEdge)(
+                short_edge_length=800,
+                max_size=1333,
+            ),
+        ],
+        augmentation_with_crop=None,
+        is_train=True,
+        mask_on=False,
+        img_format="RGB",
+    ),
+    num_workers=4,
+)
+
+# using LVIS evaluator
+dataloader.evaluator = L(LVISEvaluator)(
+    dataset_name="${..test.dataset.names}",
+)
\ No newline at end of file