remove background

Author: MateoLostanlen

dvc.yaml

@@ -20,6 +20,7 @@ stages:
       --output-dir ./data/01_model_input/yolo_train_val_small
       --sampling-ratio 1
       --random-seed 0
+      --remove-background
       --loglevel info
     - >-
       uv run python ./scripts/model/yolo/train.py

scripts/data/model_input/build.py

@@ -41,6 +41,12 @@ def make_cli_parser() -> argparse.ArgumentParser:
         default=0,
         type=int,
     )
+    parser.add_argument(
+        "--remove-background",
+        help="Remove background (negative) images that have empty label files",
+        action="store_true",
+        default=False,
+    )
     parser.add_argument(
         "-log",
         "--loglevel",
@@ -116,6 +122,7 @@ def sample_dataset(
     output_dir: Path,
     sampling_ratio: float = 0.1,
     random_seed: int = 0,
+    remove_background: bool = False,
 ) -> list[dict]:
     """
     Return a downsampled list of images and labels for the given
@@ -141,6 +148,15 @@ def sample_dataset(
                 images_filepaths,
                 k=k,
             )
+        if remove_background:
+            downsampled_image_filepaths = [
+                fp for fp in downsampled_image_filepaths
+                if (labels_split_dir / f"{fp.stem}.txt").exists()
+                and (labels_split_dir / f"{fp.stem}.txt").stat().st_size > 0
+            ]
+            logging.info(
+                f"[{split}] Removed background images. Remaining: {len(downsampled_image_filepaths)}"
+            )
         downsampled_label_filepaths = [
             labels_split_dir / f"{fp.stem}.txt"
             for fp in downsampled_image_filepaths
@@ -190,6 +206,7 @@ def run_file_copy(copy_data: list[dict]) -> None:
         output_dir = args["output_dir"]
         sampling_ratio = args["sampling_ratio"]
         random_seed = args["random_seed"]
+        remove_background = args["remove_background"]
 
         logging.info(f"Creating dirs at {output_dir}")
         shutil.rmtree(output_dir, ignore_errors=True)
@@ -202,6 +219,7 @@ def run_file_copy(copy_data: list[dict]) -> None:
                 output_dir=output_dir / "datasets",
                 sampling_ratio=sampling_ratio,
                 random_seed=random_seed,
+                remove_background=remove_background,
             )
         )
         write_data_yaml(output_dir / "datasets" / "data.yaml")