From d4340450b7ef0632dfb4b492d5dd1a6ee29ba53d Mon Sep 17 00:00:00 2001 From: Harshini Pulivarthi Date: Mon, 8 Jun 2026 16:15:23 -0700 Subject: [PATCH 1/4] modified generate.py --- .gitignore | 8 +++++++- generate.py | 9 ++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 7b95cc9..4b651fd 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,10 @@ test_data/ __pycache__/ -visualizer.py \ No newline at end of file +visualizer.py + +.DS_Store + +targets_fake +targets_original/ +output_yolo/ \ No newline at end of file diff --git a/generate.py b/generate.py index 2c5c4bd..e6d9b69 100644 --- a/generate.py +++ b/generate.py @@ -1,7 +1,8 @@ import os import argparse from PIL import Image -from compositor.overlay.overlay_multiple import paste_multiple_overlay +from compositor.overlay.overlay_single import paste_single_overlay + def load_images_from_folder(folder): files = [] @@ -9,7 +10,9 @@ def load_images_from_folder(folder): if f.lower().endswith((".png", ".jpg", ".jpeg")): path = os.path.join(folder, f) try: - files.append(Image.open(path)) + img = Image.open(path) + files.append(img.copy()) + img.close() except: print(f"Could not load image: {path}") return files @@ -29,7 +32,7 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou backgrounds = backgrounds[:num_backgrounds] for i, bg in enumerate(backgrounds): - composite, labels = paste_multiple_overlay( + composite, labels = paste_single_overlay( bg, real_targets, fake_targets, max_attempts=max_attempts ) From ddcf651e6960b31093bdb9ab09a733646a9c71ad Mon Sep 17 00:00:00 2001 From: Harshini Pulivarthi Date: Thu, 11 Jun 2026 17:20:55 -0700 Subject: [PATCH 2/4] ci: add tests, linting, and GitHub Actions workflow --- .github/workflows/ci.yml | 35 +++ README.md | 23 ++ compositor/augment.py | 5 +- compositor/overlay/overlay_multiple.py | 43 +-- compositor/overlay/overlay_single.py | 34 ++- compositor/overlay/overlay_torpedo.py | 7 +- conftest.py | 0 generate.py | 41 ++- generate_multiple.py | 32 ++- pytest.ini | 2 + requirements.txt | 3 + tests/fixtures/backgrounds/bg1.jpg | Bin 0 -> 5428 bytes tests/fixtures/backgrounds/bg2.jpg | Bin 0 -> 5429 bytes tests/fixtures/targets/target1.png | Bin 0 -> 210 bytes tests/fixtures/targets/target2.png | Bin 0 -> 209 bytes tests/test_generate.py | 82 ++++++ tests/test_overlay_single.py | 357 +++++++++++++++++++++++++ tests/test_smoke.py | 124 +++++++++ visualize_labels.py | 34 +-- 19 files changed, 758 insertions(+), 64 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 conftest.py create mode 100644 pytest.ini create mode 100644 tests/fixtures/backgrounds/bg1.jpg create mode 100644 tests/fixtures/backgrounds/bg2.jpg create mode 100644 tests/fixtures/targets/target1.png create mode 100644 tests/fixtures/targets/target2.png create mode 100644 tests/test_generate.py create mode 100644 tests/test_overlay_single.py create mode 100644 tests/test_smoke.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..3020e63 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,35 @@ +name: CI + +on: + push: + branches: ["**"] + pull_request: + branches: ["**"] + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install system deps (OpenCV) + run: sudo apt-get install -y libgl1 + + - name: Install dependencies + run: | + pip install -r requirements.txt + pip install pytest ruff black + + - name: Lint (ruff) + run: ruff check . + + - name: Format check (black) + run: black --check . + + - name: Run tests + run: pytest tests/ -v --tb=short \ No newline at end of file diff --git a/README.md b/README.md index 6166369..2aa8eb1 100644 --- a/README.md +++ b/README.md @@ -142,3 +142,26 @@ python generate.py \ ```bash python3 visualize_labels.py ``` +## 5. Testing +In order to test if generated data produces the right boxes we use the following files to test our outputs: +- tests/test_generate.py — unit tests for path loading and label file writing in the main pipeline +- tests/test_overlay_single.py — unit tests for the overlay logic including bbox math, normalization, color matching, and regression detection +- tests/test_smoke.py — end-to-end pipeline test using real fixture images to validate full output correctness +- tests/fixtures/ — tiny sample backgrounds and targets used by the smoke test +- pytest.ini — pytest configuration telling it to treat the repo root as the Python path +- conftest.py — empty file that marks the repo root as the pytest project root +- .github/workflows/ci.yml — GitHub Actions workflow that runs linting, formatting checks, and all tests on every push and pull request + +How to run tests individually or by file: + +``` bash +pytest tests/ -v +``` + +Use the following command to run all of the tests at once: + +``` bash +pytest tests/ -v +``` + +If successful it should say "30 passed in 2.04s," there are 30 tests total. diff --git a/compositor/augment.py b/compositor/augment.py index d2baea6..d09d2f6 100644 --- a/compositor/augment.py +++ b/compositor/augment.py @@ -3,6 +3,7 @@ import numpy as np from PIL import Image, ImageEnhance + def augment( image: Image.Image, brightness_range: tuple = (0.5, 1.2), @@ -30,7 +31,9 @@ def augment( # Very slight Gaussian Blur if random.random() < blur_prob: - k = random.choice(range(blur_kernel_range[0], blur_kernel_range[1] + 1, 2)) # must be odd + k = random.choice( + range(blur_kernel_range[0], blur_kernel_range[1] + 1, 2) + ) # must be odd img_np = cv2.GaussianBlur(img_np, (k, k), 0) return Image.fromarray(img_np) diff --git a/compositor/overlay/overlay_multiple.py b/compositor/overlay/overlay_multiple.py index 9aed401..536b8c4 100644 --- a/compositor/overlay/overlay_multiple.py +++ b/compositor/overlay/overlay_multiple.py @@ -3,7 +3,6 @@ import numpy as np from PIL import Image import cv2 -import os def random_perspective_transform(img_np, max_offset=0.1): @@ -12,15 +11,19 @@ def random_perspective_transform(img_np, max_offset=0.1): dy = int(h * max_offset) src = np.float32([[0, 0], [w, 0], [w, h], [0, h]]) - dst = np.float32([ - [random.randint(0, dx), random.randint(0, dy)], - [w - random.randint(0, dx), random.randint(0, dy)], - [w - random.randint(0, dx), h - random.randint(0, dy)], - [random.randint(0, dx), h - random.randint(0, dy)] - ]) + dst = np.float32( + [ + [random.randint(0, dx), random.randint(0, dy)], + [w - random.randint(0, dx), random.randint(0, dy)], + [w - random.randint(0, dx), h - random.randint(0, dy)], + [random.randint(0, dx), h - random.randint(0, dy)], + ] + ) M = cv2.getPerspectiveTransform(src, dst) - warped = cv2.warpPerspective(img_np, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0,0)) + warped = cv2.warpPerspective( + img_np, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0) + ) return warped @@ -33,8 +36,8 @@ def match_hue_subtle(overlay_np, background_np): ov_rgb = overlay_np[..., :3].astype(np.float32) bg_rgb = background_np[..., :3].astype(np.float32) - ov_mean = ov_rgb.mean(axis=(0,1)) - bg_mean = bg_rgb.mean(axis=(0,1)) + ov_mean = ov_rgb.mean(axis=(0, 1)) + bg_mean = bg_rgb.mean(axis=(0, 1)) ov_rgb += 0.05 * (bg_mean - ov_mean) ov_rgb = np.clip(ov_rgb, 0, 255).astype(np.uint8) @@ -107,12 +110,17 @@ def paste_multiple_overlay(background, real_overlays, fake_overlays, max_attempt new_box = (x_min, y_min, new_w, new_h) if not check_overlap(new_box, placed_boxes): - alpha = ov_np[..., 3] / 255.0 if ov_np.shape[2] == 4 else np.ones((new_h, new_w)) + alpha = ( + ov_np[..., 3] / 255.0 + if ov_np.shape[2] == 4 + else np.ones((new_h, new_w)) + ) for c in range(3): - bg_np[y_min:y_min+new_h, x_min:x_min+new_w, c] = ( - alpha * ov_np[..., c] + - (1 - alpha) * bg_np[y_min:y_min+new_h, x_min:x_min+new_w, c] + bg_np[y_min : y_min + new_h, x_min : x_min + new_w, c] = ( + alpha * ov_np[..., c] + + (1 - alpha) + * bg_np[y_min : y_min + new_h, x_min : x_min + new_w, c] ) placed_boxes.append(new_box) @@ -122,10 +130,9 @@ def paste_multiple_overlay(background, real_overlays, fake_overlays, max_attempt w_norm = new_w / bg.width h_norm = new_h / bg.height - labels.append({ - "class_id": class_id, - "bbox": (x_center, y_center, w_norm, h_norm) - }) + labels.append( + {"class_id": class_id, "bbox": (x_center, y_center, w_norm, h_norm)} + ) success = True break diff --git a/compositor/overlay/overlay_single.py b/compositor/overlay/overlay_single.py index b35629d..94bd487 100644 --- a/compositor/overlay/overlay_single.py +++ b/compositor/overlay/overlay_single.py @@ -3,24 +3,36 @@ from PIL import Image import cv2 + def random_perspective_transform(img_np, max_offset=0.1): h, w = img_np.shape[:2] dx = int(w * max_offset) dy = int(h * max_offset) src = np.float32([[0, 0], [w, 0], [w, h], [0, h]]) - dst = np.float32([ - [random.randint(0, dx), random.randint(0, dy)], - [w - random.randint(0, dx), random.randint(0, dy)], - [w - random.randint(0, dx), h - random.randint(0, dy)], - [random.randint(0, dx), h - random.randint(0, dy)] - ]) + dst = np.float32( + [ + [random.randint(0, dx), random.randint(0, dy)], + [w - random.randint(0, dx), random.randint(0, dy)], + [w - random.randint(0, dx), h - random.randint(0, dy)], + [random.randint(0, dx), h - random.randint(0, dy)], + ] + ) M = cv2.getPerspectiveTransform(src, dst) - warped = cv2.warpPerspective(img_np, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0)) + warped = cv2.warpPerspective( + img_np, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0) + ) return warped -def paste_single_overlay(background: Image.Image, overlay: Image.Image, class_id: int, - scale_range=(0.2, 0.6), max_attempts=20, color_match_strength=0.05): + +def paste_single_overlay( + background: Image.Image, + overlay: Image.Image, + class_id: int, + scale_range=(0.2, 0.6), + max_attempts=20, + color_match_strength=0.05, +): """ Paste a single target image onto a background with: - slight perspective warp @@ -53,7 +65,9 @@ def paste_single_overlay(background: Image.Image, overlay: Image.Image, class_id y_min = random.randint(0, bg_h - ov.height) # --- Subtle color matching --- - bg_crop = bg.crop((x_min, y_min, x_min + ov.width, y_min + ov.height)).convert("RGB") + bg_crop = bg.crop((x_min, y_min, x_min + ov.width, y_min + ov.height)).convert( + "RGB" + ) ov_rgb = ov.convert("RGB") bg_np = np.array(bg_crop).astype(np.float32) diff --git a/compositor/overlay/overlay_torpedo.py b/compositor/overlay/overlay_torpedo.py index bcce835..58f3ac3 100644 --- a/compositor/overlay/overlay_torpedo.py +++ b/compositor/overlay/overlay_torpedo.py @@ -62,10 +62,9 @@ def main(images_dir, yolo_dir, output_dir, output_yolo_dir): box_x_px = x * w_img box_y_px = y * h_img - left = box_x_px - box_w_px / 2 - top = box_y_px - box_h_px / 2 - right = box_x_px + box_w_px / 2 - bottom = box_y_px + box_h_px / 2 + left = box_x_px - box_w_px / 2 + top = box_y_px - box_h_px / 2 + right = box_x_px + box_w_px / 2 ring_center_x = (left + right) / 2 ring_center_y = top + box_h_px * 0.75 diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..e69de29 diff --git a/generate.py b/generate.py index 6f28d12..03fe36e 100644 --- a/generate.py +++ b/generate.py @@ -4,16 +4,28 @@ import random from compositor.overlay.overlay_single import paste_single_overlay -#LazyLoading: Loads paths instead of all the images at once so that the system doesn't crash - # useful because strings hold less memory than + +# LazyLoading: Loads paths instead of all the images at once so that the system doesn't crash +# useful because strings hold less memory than def load_paths_from_folder(folder): paths = [] for f in os.listdir(folder): - if f.lower().endswith((".png", ".jpg", ".jpeg")): - paths.append(os.path.join(folder, f)) #attached any file path that is an image within the folder. + if f.lower().endswith((".png", ".jpg", ".jpeg")): + paths.append( + os.path.join(folder, f) + ) # attached any file path that is an image within the folder. return paths -def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, output_yolo_dir, max_attempts, num_backgrounds=None): + +def main( + backgrounds_dir, + real_targets_dir, + fake_targets_dir, + output_img_dir, + output_yolo_dir, + max_attempts, + num_backgrounds=None, +): os.makedirs(output_img_dir, exist_ok=True) os.makedirs(output_yolo_dir, exist_ok=True) @@ -24,15 +36,18 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou background_paths = background_paths[:num_backgrounds] for i, bg_path in enumerate(background_paths): - #chooses a random index from the target images to use on each background + # chooses a random index from the target images to use on each background target_idx = random.randint(0, len(real_target_paths) - 1) class_id = target_idx + 1 - #Adjusted opening method to use paths instead of the actual image. + # Adjusted opening method to use paths instead of the actual image. with Image.open(bg_path) as bg: with Image.open(real_target_paths[target_idx]) as target: composite, labels = paste_single_overlay( - bg.copy(), target.copy(), class_id=class_id, max_attempts=max_attempts + bg.copy(), + target.copy(), + class_id=class_id, + max_attempts=max_attempts, ) img_name = f"img_{i:05d}.jpg" @@ -41,6 +56,11 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou img_path = os.path.join(output_img_dir, img_name) yolo_path = os.path.join(output_yolo_dir, label_name) + if composite is None or labels is None: + print(f"Warning: placement failed for {img_name}, saving empty label") + open(yolo_path, "w").close() + continue + composite.save(img_path, quality=95) bbox = labels.get("bbox") @@ -55,6 +75,7 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou print(f"Generated {img_name} + {label_name}") + if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -75,5 +96,5 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou args.output_img_dir, args.output_yolo_dir, args.max_attempts, - args.num_backgrounds - ) \ No newline at end of file + args.num_backgrounds, + ) diff --git a/generate_multiple.py b/generate_multiple.py index ad6a1b3..8ad29f0 100644 --- a/generate_multiple.py +++ b/generate_multiple.py @@ -3,6 +3,7 @@ from PIL import Image from compositor.overlay.overlay_multiple import paste_multiple_overlay + def load_images_from_folder(folder): files = [] for f in os.listdir(folder): @@ -10,11 +11,20 @@ def load_images_from_folder(folder): path = os.path.join(folder, f) try: files.append(Image.open(path)) - except: + except Exception: print(f"Could not load image: {path}") return files -def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, output_yolo_dir, max_attempts, num_backgrounds=None): + +def main( + backgrounds_dir, + real_targets_dir, + fake_targets_dir, + output_img_dir, + output_yolo_dir, + max_attempts, + num_backgrounds=None, +): os.makedirs(output_img_dir, exist_ok=True) os.makedirs(output_yolo_dir, exist_ok=True) @@ -49,18 +59,28 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou print(f"Generated {img_name} + {label_name}") + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--backgrounds_dir", required=True) parser.add_argument("--real_targets_dir", required=True) - parser.add_argument("--fake_targets_dir", default=None, help="Folder with fake target images (optional)") + parser.add_argument( + "--fake_targets_dir", + default=None, + help="Folder with fake target images (optional)", + ) parser.add_argument("--output_img_dir", required=True) parser.add_argument("--output_yolo_dir", required=True) parser.add_argument("--max_attempts", type=int, default=20) - parser.add_argument("--num_backgrounds", type=int, default=None, help="Number of backgrounds to process") + parser.add_argument( + "--num_backgrounds", + type=int, + default=None, + help="Number of backgrounds to process", + ) args = parser.parse_args() @@ -71,5 +91,5 @@ def main(backgrounds_dir, real_targets_dir, fake_targets_dir, output_img_dir, ou args.output_img_dir, args.output_yolo_dir, args.max_attempts, - args.num_backgrounds - ) \ No newline at end of file + args.num_backgrounds, + ) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..03f586d --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +pythonpath = . \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 55ff550..0009178 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,6 @@ numpy==2.3.0 opencv-python==4.11.0.86 pillow==11.2.1 +pytest +ruff +black \ No newline at end of file diff --git a/tests/fixtures/backgrounds/bg1.jpg b/tests/fixtures/backgrounds/bg1.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4aa7dafe525d5281a866629e8de5641af97ba6f9 GIT binary patch literal 5428 zcmex=^(PF6}rMnOeST|r4lSw=>~TvNxu(8R<| zHF0u@iAzXIsj8`KXlj|5nweWzS~We&gn?hmRgVdHU@6i$mSee*Oaai;;mD;w>Nv@fe!F1cClyVqsxs zVF&q(k*OSrnFU!`6%E;h90S=C3x$=88aYIqCNA7~kW<+>=!0ld(M2vX6_bamA3FjqiJ9?4UDFN(KIj^)4>0m0L>u?>;M1& literal 0 HcmV?d00001 diff --git a/tests/fixtures/backgrounds/bg2.jpg b/tests/fixtures/backgrounds/bg2.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c97c038e0c9f77bc46132265656933e21d6a8da6 GIT binary patch literal 5429 zcmex=^(PF6}rMnOeST|r4lSw=>~TvNxu(8R<| zHF0u@iAzXIsj8`KXlj|5nweWzS~We&gn?hmRgVdHU@6i$mSee*Oaai;;mD;w>Nv@fe!F1cClyVqsxs zVF&q(k*OSrnFU!`6%E;h90S=C3x$=88aYIqCNA7~kW<+>=!0ld(M2vX6_bamA3pj3PBWd3O?@*J=tc%lS3j3^P6Ul literal 0 HcmV?d00001 diff --git a/tests/test_generate.py b/tests/test_generate.py new file mode 100644 index 0000000..81996c8 --- /dev/null +++ b/tests/test_generate.py @@ -0,0 +1,82 @@ +from unittest.mock import patch +from PIL import Image +from generate import load_paths_from_folder, main + +# pytest tests/ -v + +# --- load_paths_from_folder --- + + +def test_load_paths_finds_images(tmp_path): + (tmp_path / "a.jpg").write_bytes(b"") + (tmp_path / "b.PNG").write_bytes(b"") + (tmp_path / "ignore.txt").write_bytes(b"") + paths = load_paths_from_folder(str(tmp_path)) + assert len(paths) == 2 + assert all(p.endswith((".jpg", ".PNG")) for p in paths) + + +def test_load_paths_empty_folder(tmp_path): + assert load_paths_from_folder(str(tmp_path)) == [] + + +# --- YOLO label format --- + + +def test_yolo_label_written_correctly(tmp_path): + """Golden test: fixed seed + mock overlay → assert exact label file content.""" + img_dir = tmp_path / "images" + lbl_dir = tmp_path / "labels" + + # Create fake background and target images + bg_dir = tmp_path / "bgs" + tgt_dir = tmp_path / "targets" + bg_dir.mkdir() + tgt_dir.mkdir() + Image.new("RGB", (640, 480), color="blue").save(bg_dir / "bg.jpg") + Image.new("RGBA", (50, 50), color="red").save(tgt_dir / "target.png") + + fake_labels = {"bbox": (0.5, 0.4, 0.1, 0.2)} + fake_composite = Image.new("RGB", (640, 480)) + + with patch( + "generate.paste_single_overlay", return_value=(fake_composite, fake_labels) + ): + main( + str(bg_dir), str(tgt_dir), None, str(img_dir), str(lbl_dir), max_attempts=5 + ) + + label_file = lbl_dir / "img_00000.txt" + assert label_file.exists() + line = label_file.read_text().strip() + parts = line.split() + assert len(parts) == 5 # class x y w h + assert parts[0] == "1" # class_id = target_idx + 1 + x, y, w, h = map(float, parts[1:]) + assert 0.0 <= x <= 1.0 + assert 0.0 <= y <= 1.0 + assert 0.0 < w <= 1.0 + assert 0.0 < h <= 1.0 + + +def test_empty_label_on_no_placement(tmp_path): + """If paste_single_overlay returns no bbox, label file should be empty.""" + img_dir = tmp_path / "images" + lbl_dir = tmp_path / "labels" + bg_dir = tmp_path / "bgs" + tgt_dir = tmp_path / "targets" + bg_dir.mkdir() + tgt_dir.mkdir() + Image.new("RGB", (640, 480)).save(bg_dir / "bg.jpg") + Image.new("RGBA", (50, 50)).save(tgt_dir / "target.png") + + with patch( + "generate.paste_single_overlay", return_value=(Image.new("RGB", (640, 480)), {}) + ): + main( + str(bg_dir), str(tgt_dir), None, str(img_dir), str(lbl_dir), max_attempts=5 + ) + + label_file = lbl_dir / "img_00000.txt" + assert label_file.exists() + assert label_file.read_text() == "" # empty label = valid YOLO negative diff --git a/tests/test_overlay_single.py b/tests/test_overlay_single.py new file mode 100644 index 0000000..76b3791 --- /dev/null +++ b/tests/test_overlay_single.py @@ -0,0 +1,357 @@ +""" +Tests for compositor/overlay/overlay_single.py + +Run with: + pytest tests/test_overlay_single.py -v +""" + +import random +import numpy as np +import pytest +from PIL import Image +from unittest.mock import patch + +from compositor.overlay.overlay_single import ( + paste_single_overlay, + random_perspective_transform, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def make_bg(w=640, h=480, color="blue"): + return Image.new("RGB", (w, h), color=color) + + +def make_target(w=80, h=80, color=(255, 0, 0, 255)): + img = Image.new("RGBA", (w, h), color=color) + return img + + +# --------------------------------------------------------------------------- +# random_perspective_transform +# --------------------------------------------------------------------------- + + +class TestPerspectiveTransform: + def test_output_shape_preserved(self): + img = np.zeros((100, 100, 4), dtype=np.uint8) + out = random_perspective_transform(img) + assert out.shape == img.shape, "Output shape must match input shape" + + def test_output_dtype_preserved(self): + img = np.zeros((100, 100, 4), dtype=np.uint8) + out = random_perspective_transform(img) + assert out.dtype == np.uint8 + + def test_no_nans_in_output(self): + img = np.random.randint(0, 255, (100, 100, 4), dtype=np.uint8) + out = random_perspective_transform(img) + assert not np.isnan( + out.astype(float) + ).any(), "Perspective transform must not produce NaNs" + + def test_deterministic_with_seed(self): + img = np.random.randint(0, 255, (100, 100, 4), dtype=np.uint8) + random.seed(42) + out1 = random_perspective_transform(img) + random.seed(42) + out2 = random_perspective_transform(img) + np.testing.assert_array_equal(out1, out2) + + def test_zero_offset_is_identity(self): + """max_offset=0 means no warp — output should equal input.""" + img = np.random.randint(0, 255, (100, 100, 4), dtype=np.uint8) + out = random_perspective_transform(img, max_offset=0.0) + np.testing.assert_array_equal(out, img) + + +# --------------------------------------------------------------------------- +# paste_single_overlay — return value contract +# --------------------------------------------------------------------------- + + +class TestPasteSingleOverlayContract: + def test_returns_image_and_label_on_success(self): + random.seed(0) + bg = make_bg() + tgt = make_target() + result_img, label = paste_single_overlay(bg, tgt, class_id=1) + assert result_img is not None, "Should return a composite image" + assert label is not None, "Should return a label dict" + + def test_returns_none_none_when_max_attempts_exhausted(self): + """Force failure by making overlay larger than 40% of bg.""" + bg = make_bg(100, 100) + # Target nearly as big as background → scale_range forces area > 40% + tgt = make_target(90, 90) + result_img, label = paste_single_overlay( + bg, tgt, class_id=1, scale_range=(0.95, 0.99), max_attempts=5 + ) + assert result_img is None + assert label is None + + def test_label_has_bbox_key(self): + random.seed(1) + bg = make_bg() + tgt = make_target() + _, label = paste_single_overlay(bg, tgt, class_id=3) + assert "bbox" in label + + def test_label_has_class_id_key(self): + random.seed(2) + bg = make_bg() + tgt = make_target() + _, label = paste_single_overlay(bg, tgt, class_id=5) + assert label["class_id"] == 5 + + def test_output_image_is_rgb(self): + random.seed(3) + bg = make_bg() + tgt = make_target() + result_img, _ = paste_single_overlay(bg, tgt, class_id=1) + assert result_img.mode == "RGB" + + def test_output_image_same_size_as_background(self): + random.seed(4) + bg = make_bg(320, 240) + tgt = make_target(40, 40) + result_img, _ = paste_single_overlay(bg, tgt, class_id=1) + assert result_img.size == (320, 240) + + +# --------------------------------------------------------------------------- +# paste_single_overlay — YOLO bbox normalization math +# --------------------------------------------------------------------------- + + +class TestYOLONormalization: + """ + These tests lock down the coordinate math. We mock random calls so the + placement is deterministic and we can compute expected values exactly. + """ + + def _run_with_fixed_placement(self, bg_w, bg_h, ov_w, ov_h, x_min, y_min): + """ + Bypass all randomness: fixed scale produces a known overlay size, + fixed position produces a known placement. + Returns (x_center, y_center, w_norm, h_norm). + """ + random.seed(0) + bg = make_bg(bg_w, bg_h) + tgt = make_target(ov_w, ov_h) + + # Patch the perspective transform to be a no-op so overlay dims don't change + def noop_transform(img_np, max_offset=0.1): + return img_np + + with patch( + "compositor.overlay.overlay_single.random_perspective_transform", + side_effect=noop_transform, + ), patch( + "compositor.overlay.overlay_single.random.uniform", return_value=ov_w / bg_w + ), patch( + "compositor.overlay.overlay_single.random.randint", + side_effect=[x_min, y_min], + ): + result_img, label = paste_single_overlay( + bg, tgt, class_id=1, max_attempts=1 + ) + + return label["bbox"] if label else None + + def test_centered_placement(self): + """Overlay placed dead-center → x_center=0.5, y_center=0.5.""" + bg_w, bg_h = 640, 480 + ov_w, ov_h = 128, 128 + x_min = (bg_w - ov_w) // 2 # 256 + y_min = (bg_h - ov_h) // 2 # 176 + + bbox = self._run_with_fixed_placement(bg_w, bg_h, ov_w, ov_h, x_min, y_min) + if bbox is None: + pytest.skip("Placement failed — check scale constraints") + + x_c, y_c, w_n, h_n = bbox + expected_xc = (x_min + ov_w / 2) / bg_w + expected_yc = (y_min + ov_h / 2) / bg_h + assert abs(x_c - expected_xc) < 1e-5 + assert abs(y_c - expected_yc) < 1e-5 + + def test_bbox_coords_in_unit_range(self): + """All YOLO values must be in [0, 1].""" + random.seed(99) + bg = make_bg() + tgt = make_target() + _, label = paste_single_overlay(bg, tgt, class_id=1) + if label is None: + pytest.skip("No placement") + x_c, y_c, w_n, h_n = label["bbox"] + assert 0.0 <= x_c <= 1.0, f"x_center out of range: {x_c}" + assert 0.0 <= y_c <= 1.0, f"y_center out of range: {y_c}" + assert 0.0 < w_n <= 1.0, f"width out of range: {w_n}" + assert 0.0 < h_n <= 1.0, f"height out of range: {h_n}" + + def test_no_zero_area_bbox(self): + random.seed(7) + bg = make_bg() + tgt = make_target() + _, label = paste_single_overlay(bg, tgt, class_id=1) + if label is None: + pytest.skip("No placement") + _, _, w_n, h_n = label["bbox"] + assert w_n > 0 and h_n > 0, "Bounding box must have non-zero area" + + def test_no_nan_in_bbox(self): + random.seed(8) + bg = make_bg() + tgt = make_target() + _, label = paste_single_overlay(bg, tgt, class_id=1) + if label is None: + pytest.skip("No placement") + for val in label["bbox"]: + assert not (val != val), f"NaN found in bbox: {label['bbox']}" # NaN != NaN + + def test_width_proportional_to_scale(self): + """w_norm ≈ scale (since overlay width = bg_w * scale, before perspective warp).""" + random.seed(10) + bg = make_bg(640, 480) + tgt = make_target(100, 100) + + def noop_transform(img_np, max_offset=0.1): + return img_np + + with patch( + "compositor.overlay.overlay_single.random_perspective_transform", + side_effect=noop_transform, + ): + _, label = paste_single_overlay( + bg, tgt, class_id=1, scale_range=(0.3, 0.3), max_attempts=10 + ) + + if label is None: + pytest.skip("No placement") + _, _, w_n, _ = label["bbox"] + assert abs(w_n - 0.3) < 0.02, f"Expected w_norm ≈ 0.3, got {w_n}" + + +# --------------------------------------------------------------------------- +# paste_single_overlay — area constraint +# --------------------------------------------------------------------------- + + +class TestAreaConstraint: + def test_overlay_never_exceeds_40_percent_of_background(self): + """The function must skip placements where overlay area > 40% of bg.""" + random.seed(42) + bg = make_bg(640, 480) + tgt = make_target(80, 80) + bg_area = 640 * 480 + + for seed in range(20): + random.seed(seed) + result_img, label = paste_single_overlay(bg.copy(), tgt.copy(), class_id=1) + if label is None: + continue + _, _, w_n, h_n = label["bbox"] + overlay_area = (w_n * 640) * (h_n * 480) + assert ( + overlay_area <= 0.4 * bg_area + 1 + ), f"Overlay area {overlay_area:.0f} exceeds 40% of bg ({0.4 * bg_area:.0f})" + + +# --------------------------------------------------------------------------- +# paste_single_overlay — color matching robustness +# --------------------------------------------------------------------------- + + +class TestColorMatching: + def test_color_match_strength_zero_leaves_overlay_unchanged(self): + """With color_match_strength=0, overlay RGB should not shift at all.""" + random.seed(5) + bg = make_bg(640, 480, color="black") + # Bright red target — would shift strongly with nonzero strength + tgt = Image.new("RGBA", (80, 80), color=(255, 0, 0, 255)) + + def noop_transform(img_np, max_offset=0.1): + return img_np + + with patch( + "compositor.overlay.overlay_single.random_perspective_transform", + side_effect=noop_transform, + ): + result_img, label = paste_single_overlay( + bg, tgt, class_id=1, color_match_strength=0.0 + ) + + assert result_img is not None + + def test_output_has_no_nans(self): + """Composite image pixel values must not contain NaN.""" + random.seed(6) + bg = make_bg() + tgt = make_target() + result_img, _ = paste_single_overlay(bg, tgt, class_id=1) + if result_img is None: + pytest.skip("No placement") + arr = np.array(result_img, dtype=np.float32) + assert not np.isnan(arr).any() + + def test_output_pixels_in_valid_range(self): + """All pixel values must be 0–255.""" + random.seed(11) + bg = make_bg() + tgt = make_target() + result_img, _ = paste_single_overlay(bg, tgt, class_id=1) + if result_img is None: + pytest.skip("No placement") + arr = np.array(result_img) + assert arr.min() >= 0 and arr.max() <= 255 + + +# --------------------------------------------------------------------------- +# Golden test — regression anchor +# --------------------------------------------------------------------------- + + +class TestGoldenRegression: + """ + Fixed seed + fixed inputs → fixed bbox. If this breaks, something subtle + changed in the placement or normalization logic. + """ + + GOLDEN_SEED = 42 + EXPECTED_BBOX = ( + 0.296094, + 0.677083, + 0.454688, + 0.454167, + ) # found using pytest -v -k test_golden_bbox_is_stable -s + + def test_golden_bbox_is_stable(self): + """ + HOW TO USE: + 1. Run once: pytest -v -k test_golden_bbox_is_stable -s + 2. Copy the printed bbox into EXPECTED_BBOX above. + 3. From then on this test acts as a regression guard. + """ + random.seed(self.GOLDEN_SEED) + np.random.seed(self.GOLDEN_SEED) + + bg = Image.new("RGB", (640, 480), color=(100, 149, 237)) # cornflower blue + tgt = Image.new("RGBA", (80, 60), color=(200, 50, 50, 255)) + + result_img, label = paste_single_overlay(bg, tgt, class_id=2, max_attempts=20) + + assert ( + result_img is not None + ), "Golden test: placement must succeed with seed 42" + bbox = label["bbox"] + print(f"\n[golden] bbox = {tuple(round(v, 6) for v in bbox)}") + + if self.EXPECTED_BBOX is not None: + for got, exp in zip(bbox, self.EXPECTED_BBOX): + assert ( + abs(got - exp) < 1e-4 + ), f"Regression: bbox changed. got={bbox}, expected={self.EXPECTED_BBOX}" diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..8f9c517 --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,124 @@ +""" +Smoke test for the full data generation pipeline. + +Uses real fixture images (tiny) to run the pipeline end-to-end and validate outputs. + +Run with: + pytest tests/test_smoke.py -v +""" + +import os +import pytest +from PIL import Image +from generate import main + +FIXTURES_DIR = os.path.join(os.path.dirname(__file__), "fixtures") +BACKGROUNDS_DIR = os.path.join(FIXTURES_DIR, "backgrounds") +TARGETS_DIR = os.path.join(FIXTURES_DIR, "targets") + + +@pytest.fixture +def output_dirs(tmp_path): + img_dir = tmp_path / "images" + lbl_dir = tmp_path / "labels" + img_dir.mkdir() + lbl_dir.mkdir() + return str(img_dir), str(lbl_dir) + + +class TestSmokePipeline: + def test_pipeline_runs_without_crashing(self, output_dirs): + """Full pipeline completes on real fixture images without raising.""" + img_dir, lbl_dir = output_dirs + main( + backgrounds_dir=BACKGROUNDS_DIR, + real_targets_dir=TARGETS_DIR, + fake_targets_dir=None, + output_img_dir=img_dir, + output_yolo_dir=lbl_dir, + max_attempts=20, + ) + + def test_correct_number_of_outputs(self, output_dirs): + """One image + one label file per background.""" + img_dir, lbl_dir = output_dirs + main( + backgrounds_dir=BACKGROUNDS_DIR, + real_targets_dir=TARGETS_DIR, + fake_targets_dir=None, + output_img_dir=img_dir, + output_yolo_dir=lbl_dir, + max_attempts=20, + ) + images = [f for f in os.listdir(img_dir) if f.endswith(".jpg")] + labels = [f for f in os.listdir(lbl_dir) if f.endswith(".txt")] + num_backgrounds = len([f for f in os.listdir(BACKGROUNDS_DIR)]) + assert len(images) == num_backgrounds + assert len(labels) == num_backgrounds + + def test_output_images_are_valid(self, output_dirs): + """Every output image can be opened and is RGB.""" + img_dir, lbl_dir = output_dirs + main( + backgrounds_dir=BACKGROUNDS_DIR, + real_targets_dir=TARGETS_DIR, + fake_targets_dir=None, + output_img_dir=img_dir, + output_yolo_dir=lbl_dir, + max_attempts=20, + ) + for fname in os.listdir(img_dir): + if fname.endswith(".jpg"): + img = Image.open(os.path.join(img_dir, fname)) + assert img.mode == "RGB", f"{fname} is not RGB" + assert img.size == (640, 480), f"{fname} has wrong size" + + def test_label_files_have_valid_yolo_format(self, output_dirs): + """Every non-empty label file contains valid YOLO formatted lines.""" + img_dir, lbl_dir = output_dirs + main( + backgrounds_dir=BACKGROUNDS_DIR, + real_targets_dir=TARGETS_DIR, + fake_targets_dir=None, + output_img_dir=img_dir, + output_yolo_dir=lbl_dir, + max_attempts=20, + ) + for fname in os.listdir(lbl_dir): + fpath = os.path.join(lbl_dir, fname) + content = open(fpath).read().strip() + if not content: + continue # empty label is valid (no placement) + + for line in content.splitlines(): + parts = line.strip().split() + assert len(parts) == 5, f"{fname}: expected 5 values, got {len(parts)}" + + class_id = int(parts[0]) + x, y, w, h = map(float, parts[1:]) + + assert class_id >= 1, f"{fname}: class_id must be >= 1" + assert 0.0 <= x <= 1.0, f"{fname}: x_center out of range: {x}" + assert 0.0 <= y <= 1.0, f"{fname}: y_center out of range: {y}" + assert 0.0 < w <= 1.0, f"{fname}: width out of range: {w}" + assert 0.0 < h <= 1.0, f"{fname}: height out of range: {h}" + assert w > 0 and h > 0, f"{fname}: zero-area bbox" + + # No NaNs + for val in [x, y, w, h]: + assert val == val, f"{fname}: NaN found in bbox" + + def test_num_backgrounds_limit(self, output_dirs): + """num_backgrounds parameter correctly limits output count.""" + img_dir, lbl_dir = output_dirs + main( + backgrounds_dir=BACKGROUNDS_DIR, + real_targets_dir=TARGETS_DIR, + fake_targets_dir=None, + output_img_dir=img_dir, + output_yolo_dir=lbl_dir, + max_attempts=20, + num_backgrounds=1, + ) + images = [f for f in os.listdir(img_dir) if f.endswith(".jpg")] + assert len(images) == 1 \ No newline at end of file diff --git a/visualize_labels.py b/visualize_labels.py index a4463d6..153781b 100644 --- a/visualize_labels.py +++ b/visualize_labels.py @@ -1,24 +1,24 @@ -import os import cv2 import numpy as np from pathlib import Path -IMAGE_DIR = "output_images" -LABEL_DIR = "output_yolo" -OUTPUT_DIR = "visualized" +IMAGE_DIR = "output_images" +LABEL_DIR = "output_yolo" +OUTPUT_DIR = "visualized" # Optional: map class IDs to names, e.g. {0: "cat", 1: "dog"} CLASS_NAMES = {} -#Drawing Settings -BOX_THICKNESS = 2 -FONT_SCALE = 0.6 +# Drawing Settings +BOX_THICKNESS = 2 +FONT_SCALE = 0.6 FONT_THICKNESS = 1 # If True, each class gets a unique color; if False, all boxes are green COLORIZE_BY_CLASS = True # ───────────────────────────────────────────────────────────────────────────── + # color per class/image def class_color(class_id: int): np.random.seed(class_id + 42) @@ -40,7 +40,7 @@ def draw_yolo_boxes(image_path: Path, label_path: Path, output_path: Path): return with open(label_path) as f: - lines = [l.strip() for l in f if l.strip()] + lines = [line.strip() for line in f if line.strip()] for line in lines: parts = line.split() @@ -66,18 +66,22 @@ def draw_yolo_boxes(image_path: Path, label_path: Path, output_path: Path): # Background rectangle for text cv2.rectangle(img, (x1, y1 - th - baseline - 4), (x1 + tw + 4, y1), color, -1) cv2.putText( - img, label, + img, + label, (x1 + 2, y1 - baseline - 2), - cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, - (255, 255, 255), FONT_THICKNESS, cv2.LINE_AA + cv2.FONT_HERSHEY_SIMPLEX, + FONT_SCALE, + (255, 255, 255), + FONT_THICKNESS, + cv2.LINE_AA, ) cv2.imwrite(str(output_path), img) def main(): - image_dir = Path(IMAGE_DIR) - label_dir = Path(LABEL_DIR) + image_dir = Path(IMAGE_DIR) + label_dir = Path(LABEL_DIR) output_dir = Path(OUTPUT_DIR) output_dir.mkdir(parents=True, exist_ok=True) @@ -91,7 +95,7 @@ def main(): print(f"Found {len(image_files)} images. Processing...") for i, img_path in enumerate(sorted(image_files), 1): - label_path = label_dir / (img_path.stem + ".txt") + label_path = label_dir / (img_path.stem + ".txt") output_path = output_dir / img_path.name draw_yolo_boxes(img_path, label_path, output_path) if i % 500 == 0 or i == len(image_files): @@ -101,4 +105,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() From 3af4f66f16c6e1ee1d1fb8a57847f0c515d15dc0 Mon Sep 17 00:00:00 2001 From: Harshini Pulivarthi Date: Thu, 11 Jun 2026 17:25:51 -0700 Subject: [PATCH 3/4] modified readme --- README.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 2aa8eb1..e8282e8 100644 --- a/README.md +++ b/README.md @@ -152,10 +152,14 @@ In order to test if generated data produces the right boxes we use the following - conftest.py — empty file that marks the repo root as the pytest project root - .github/workflows/ci.yml — GitHub Actions workflow that runs linting, formatting checks, and all tests on every push and pull request -How to run tests individually or by file: +Running tests by file: ``` bash -pytest tests/ -v +pytest tests/test_generate.py -v +pytest tests/test_overlay_single.py -v +pytest tests/test_smoke.py -v + + ``` Use the following command to run all of the tests at once: From cfabb9311bd39901db32605ff77c34489b03af7f Mon Sep 17 00:00:00 2001 From: Harshini Pulivarthi Date: Thu, 11 Jun 2026 18:39:46 -0700 Subject: [PATCH 4/4] style: black format test_smoke.py --- tests/test_smoke.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 8f9c517..1783f3a 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -100,9 +100,9 @@ def test_label_files_have_valid_yolo_format(self, output_dirs): assert class_id >= 1, f"{fname}: class_id must be >= 1" assert 0.0 <= x <= 1.0, f"{fname}: x_center out of range: {x}" assert 0.0 <= y <= 1.0, f"{fname}: y_center out of range: {y}" - assert 0.0 < w <= 1.0, f"{fname}: width out of range: {w}" - assert 0.0 < h <= 1.0, f"{fname}: height out of range: {h}" - assert w > 0 and h > 0, f"{fname}: zero-area bbox" + assert 0.0 < w <= 1.0, f"{fname}: width out of range: {w}" + assert 0.0 < h <= 1.0, f"{fname}: height out of range: {h}" + assert w > 0 and h > 0, f"{fname}: zero-area bbox" # No NaNs for val in [x, y, w, h]: @@ -121,4 +121,4 @@ def test_num_backgrounds_limit(self, output_dirs): num_backgrounds=1, ) images = [f for f in os.listdir(img_dir) if f.endswith(".jpg")] - assert len(images) == 1 \ No newline at end of file + assert len(images) == 1