From 7b4c899456e2f5d46139bfb842370ce556d7707c Mon Sep 17 00:00:00 2001
From: Eugene Liu <eugene.liu@intel.com>
Date: Thu, 19 Dec 2024 17:11:18 +0000
Subject: [PATCH] Refactor D-FINE module by removing commented code, cleaning
 up imports, and updating documentation

---
 src/otx/algo/detection/losses/dfine_loss.py   |  2 +
 src/otx/core/data/dataset/detection.py        |  1 -
 .../core/data/transform_libs/torchvision.py   | 38 +------------------
 src/otx/core/exporter/native.py               |  5 ---
 src/otx/core/model/base.py                    |  3 --
 src/otx/recipe/_base_/train.yaml              |  3 --
 6 files changed, 4 insertions(+), 48 deletions(-)

diff --git a/src/otx/algo/detection/losses/dfine_loss.py b/src/otx/algo/detection/losses/dfine_loss.py
index 6540c802f32..ddddc7f2f81 100644
--- a/src/otx/algo/detection/losses/dfine_loss.py
+++ b/src/otx/algo/detection/losses/dfine_loss.py
@@ -128,6 +128,8 @@ def bbox2distance(
 class DFINECriterion(nn.Module):
     """D-Fine criterion with FGL and DDF losses.
 
+    TODO(Eugene): Consider merge with RTDETRCriterion in the next PR.
+
     The process happens in two steps:
     1) we compute hungarian assignment between ground truth boxes and the outputs of the model
     2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
diff --git a/src/otx/core/data/dataset/detection.py b/src/otx/core/data/dataset/detection.py
index 1ffea8fd974..429ac77b71b 100644
--- a/src/otx/core/data/dataset/detection.py
+++ b/src/otx/core/data/dataset/detection.py
@@ -22,7 +22,6 @@
 class OTXDetectionDataset(OTXDataset[DetDataEntity]):
     """OTXDataset class for detection task."""
 
-    # @timing
     def _get_item_impl(self, index: int) -> DetDataEntity | None:
         item = self.dm_subset[index]
         img = item.media_as(Image)
diff --git a/src/otx/core/data/transform_libs/torchvision.py b/src/otx/core/data/transform_libs/torchvision.py
index f99eb32c6be..16d9c4d4532 100644
--- a/src/otx/core/data/transform_libs/torchvision.py
+++ b/src/otx/core/data/transform_libs/torchvision.py
@@ -80,42 +80,8 @@
 
 # mypy: disable-error-code="attr-defined"
 
-__all__ = [
-    "NumpytoTVTensorMixin",
-    "PerturbBoundingBoxes",
-    "PadtoSquare",
-    "ResizetoLongestEdge",
-    "DecodeVideo",
-    "PackVideo",
-    "MinIoURandomCrop",
-    "Resize",
-    "RandomResizedCrop",
-    "EfficientNetRandomCrop",
-    "RandomFlip",
-    "PhotoMetricDistortion",
-    "RandomAffine",
-    "CachedMosaic",
-    "CachedMixUp",
-    "YOLOXHSVRandomAug",
-    "Pad",
-    "RandomResize",
-    "RandomCrop",
-    "FilterAnnotations",
-    "Compose",
-    "FormatShape",
-    "DecordInit",
-    "SampleFrames",
-    "DecordDecode",
-    "Normalize3D",
-    "GetBBoxCenterScale",
-    "RandomBBoxTransform",
-    "TopdownAffine",
-    "Decode3DInputsAffineTransforms",
-    "TorchVisionTransformLib",
-]
-
-
-def custom_query_size(flat_inputs: list[Any]) -> tuple[int, int]:
+
+def custom_query_size(flat_inputs: list[Any]) -> tuple[int, int]:  # noqa: D103
     sizes = {
         tuple(F.get_size(inpt))
         for inpt in flat_inputs
diff --git a/src/otx/core/exporter/native.py b/src/otx/core/exporter/native.py
index 4bfc46f917e..5f901ede2b5 100644
--- a/src/otx/core/exporter/native.py
+++ b/src/otx/core/exporter/native.py
@@ -68,11 +68,6 @@ def to_openvino(
         """
         dummy_tensor = torch.rand(self.input_size).to(next(model.parameters()).device)
 
-        # TODO(Eugene): Benchmark the performance of the model after conversion.
-        # for m in model.modules():
-        #     if hasattr(m, 'convert_to_deploy'):
-        #         m.convert_to_deploy()
-
         if self.via_onnx:
             with tempfile.TemporaryDirectory() as tmpdirname:
                 tmp_dir = Path(tmpdirname)
diff --git a/src/otx/core/model/base.py b/src/otx/core/model/base.py
index b5b0130e201..a48325ca98c 100644
--- a/src/otx/core/model/base.py
+++ b/src/otx/core/model/base.py
@@ -148,9 +148,6 @@ def training_step(self, batch: T_OTXBatchDataEntity, batch_idx: int) -> Tensor |
             # TODO (sungchul): check this in distributed training
             return None if self.trainer.world_size == 1 else torch.tensor(0.0, device=self.device)
 
-        # if batch_idx == 0:
-        #     self.vis(batch.images[0], batch.bboxes[0])
-
         if isinstance(train_loss, Tensor):
             self.log(
                 "train/loss",
diff --git a/src/otx/recipe/_base_/train.yaml b/src/otx/recipe/_base_/train.yaml
index b501e2edc56..7dba87f8381 100644
--- a/src/otx/recipe/_base_/train.yaml
+++ b/src/otx/recipe/_base_/train.yaml
@@ -53,9 +53,6 @@ logger:
       log_graph: false
       default_hp_metric: true
       prefix: ""
-  # - class_path: lightning.pytorch.loggers.wandb.WandbLogger
-  #   init_args:
-  #     project: "dfine"
 deterministic: false
 seed: null
 precision: 16