🔀 [AAAI|Merge] branch 'feature/use-preprocess-target'

Author: henrytsui000

yolo/aaai.py

@@ -31,6 +31,12 @@ def __init__(self, cfg: Config, model):
         self.contrastive_loss = NT_Xent
         self.cfg = cfg
         self.metric = meanBoxCoverScore()
+        self.target_source_train = torch.load(
+            f"{cfg.dataset.path}/target/from_{cfg.task.data.target_source}_train.pt", weights_only=False
+        )
+        self.target_source_val = torch.load(
+            f"{cfg.dataset.path}/target/from_{cfg.task.data.target_source}_val.pt", weights_only=False
+        )
 
     def set_task(self, task):
         self.task = task
@@ -42,12 +48,14 @@ def setup(self, stage):
         )
         self.loss_fn = AAAILoss(self.cfg.task.loss, self.vec2box)
         self.post_process = PostProcess(self.vec2box, self.cfg.task.validation.nms, aaai=True)
+        self.target_source_train = self.target_source_train.to(self.device)
+        self.target_source_val = self.target_source_val.to(self.device)
 
     def forward(self, x, external=None, shortcut=None):
         return self.model(x, external, shortcut)
 
     def training_step(self, batch, batch_idx):
-        images_batch, masks_batch, puzzles_batch, idx_batch = batch
+        images_batch, masks_batch, _, idx_batch = batch
         R_loss, C_loss, D_loss = 0, 0, 0
         images, bbox = images_batch
         batch_size = images.size(0)
@@ -61,18 +69,9 @@ def training_step(self, batch, batch_idx):
             R_loss_msk = self.construct_loss(masked_outputs["RMAP"] * masks, images * masks)
             R_loss = R_loss_all + R_loss_msk * 10
 
-        if self.task == "puzzle" or self.task == "detect":
-            puzzle_images, origin_idx, puzzle_idx = puzzles_batch
-            puzzle_outputs = self(puzzle_images, shortcut="FMAP")
-            picked_puzzle = puzzle_outputs["FMAP"].detach()[batch_step[:, None], :, puzzle_idx[:, 0], puzzle_idx[:, 1]]
-
-        if self.task == "puzzle":
-            picked_origin = origin_outputs["FMAP"][:, :, origin_idx[0, 0], origin_idx[0, 1]]
-            C_loss = self.contrastive_loss(picked_puzzle, picked_origin)
-
         if self.task == "detect":
-            _, pick_idx = idx_batch
-            picked_vector = picked_puzzle.view(batch_size, -1, 512)[batch_step, pick_idx]
+            image_idx, pick_idx = idx_batch
+            picked_vector = self.target_source_train[image_idx[:, None], pick_idx]
 
             origin_outputs = self(images, dict(target=picked_vector.permute(0, 2, 1)))
             detections = self.vec2box(origin_outputs["Main"])
@@ -132,12 +131,9 @@ def validation_step(self, batch, batch_idx):
 
         if self.task == "detect":
 
-            puzzle_images, origin_idx, puzzle_idx = puzzles_batch
-            puzzle_outputs = self(puzzle_images, shortcut="FMAP")
-            picked_puzzle = puzzle_outputs["FMAP"].detach()[batch_step[:, None], :, puzzle_idx[:, 0], puzzle_idx[:, 1]]
-
-            _, pick_idx = idx_batch
-            picked_vector = picked_puzzle.view(batch_size, -1, 512)[batch_step, pick_idx]
+            puzzle_images, origin_idx, puzzle_idx, puzzles = puzzles_batch
+            image_idx, pick_idx = idx_batch
+            picked_vector = self.target_source_val[image_idx[:, None], pick_idx]
 
             origin_outputs = self(images, dict(target=picked_vector.permute(0, 2, 1)))
             H, W = images.shape[2:]
@@ -162,16 +158,13 @@ def validation_step(self, batch, batch_idx):
     def on_validation_batch_end(self, outputs, batch, batch_idx):
         if batch_idx != 0:
             return
-        images_batch, _, puzzles_batch, _ = batch
+        images_batch, _, _, _ = batch
         images, bbox = images_batch
-        puzzle_images, _, _ = puzzles_batch
         origin_image = draw_bboxes(images[0], bbox[0])
         predict_image = draw_bboxes(images[0], outputs[0])
-        puzzle_image = to_pil_image(puzzle_images[0])
         for logger in self.loggers:
             if isinstance(logger, WandbLogger):
                 logger.log_image(f"Origin Image", [origin_image], self.current_epoch)
-                logger.log_image(f"Puzzle Visualize", [puzzle_image], self.current_epoch)
                 logger.log_image(f"Predict Visualize", [predict_image], self.current_epoch)
 
     def on_validation_start(self):
@@ -236,7 +229,7 @@ def main(cfg: Config):
         deterministic=True,
         logger=loggers,
         devices=[0],
-        gradient_clip_algorithm='norm',
+        gradient_clip_algorithm="norm",
         gradient_clip_val=10,
         callbacks=[checkpoint_callback, RichProgressBar(), YOLORichModelSummary()],
         accelerator="auto",

yolo/config/config.py

@@ -66,7 +66,7 @@ class DataConfig:
 class OriginConfig:
     scale_origin: float
     background_color: int
-    target_file: str
+    target_postfix: str
 
 
 @dataclass
@@ -87,6 +87,7 @@ class AAAIDataConfig(DataConfig):
     mask: MaskConfig
     puzzle: PuzzleConfig
     num_target: int
+    target_source: str
 
 
 @dataclass

yolo/config/task/aaai.yaml

@@ -4,7 +4,7 @@ target: detect
 defaults:
   - validation: ../validation
 
-epoch: 500
+epoch: 25
 
 validation:
   nms:
@@ -22,6 +22,7 @@ data:
     resolution: 32
     size: 160
   num_target: 3
+  target_source: clip
   batch_size: 16
   image_size: ${image_size}
   cpu_num: ${cpu_num}

yolo/tools/data_loader.py

@@ -385,6 +385,14 @@ def __len__(self):
         return len(self.image_paths)
 
     def __getitem__(self, idx: int):
+        """
+        Returns:
+            - origin_image [3, H, W], bboxes [Box, num_target]
+            - masked_image [3, H, W], mask [3, H, W]
+            - puzzle_image [3, H, W], origin_grid [2, p_num, p_num], \
+                puzzle_grid [2, p_num, p_num], puzzles [3, p_num, p_num, p_size, p_size]
+            - idx [int], pick_idx [p_num]
+        """
         image_path = self.dataset_path / self.image_paths[idx]
         image = Image.open(image_path).convert("RGB")
         if self.transform:
@@ -398,10 +406,15 @@ def __getitem__(self, idx: int):
         if hasattr(self, "filter_box"):
             bboxes = bboxes[pick_idx]
         image = to_tensor(image)
-        masked_image, mask = self.mask(image)
-        puzzle_image, origin_grid, puzzle_grid = self.puzzle(image, shift_hw)
-
-        return (origin_image, bboxes), (masked_image, mask), (puzzle_image, origin_grid, puzzle_grid), (idx, pick_idx)
+        # masked_image, mask = self.mask(image)
+        # puzzle_image, origin_grid, puzzle_grid, puzzles = self.puzzle(image, shift_hw)
+
+        return (
+            (origin_image, bboxes),
+            ([], []),
+            ([], [], [], []),
+            (idx, pick_idx),
+        )
 
     def augment_origin(self, image: Image, shift_hw) -> tuple[Tensor, Tensor]:
         w, h = self.image_size[0] // self.puzzle_size, self.image_size[1] // self.puzzle_size
@@ -416,8 +429,8 @@ def augment_origin(self, image: Image, shift_hw) -> tuple[Tensor, Tensor]:
             .view(-1, 4)
             .float()
         )
-        if self.task == "detect":
-            image, bboxes = random_resize_crop(image, bboxes, self.data_cfg.main)
+        # if self.task == "detect":
+        #     image, bboxes = random_resize_crop(image, bboxes, self.data_cfg.main)
         return to_tensor(image), bboxes
 
     def puzzle(self, image: Tensor, hw: Optional[Tensor] = None) -> tuple[Tensor]: