pupil_tracking/check_augmentation.py at main · Aitchel17/pupil_tracking · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# -*- coding: utf-8 -*-
"""
Created on Thu Jan 22 17:18:21 2026

@author: yzhao
"""

from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt

from dataset import PupilDataset


def show_augmented_samples(
    dataset,
    n_samples=5,
    n_augs_per_sample=2,
    overlay_mask=True,
    mask_transparency=0.1,
):
    """
    Visualize augmented samples from a PupilDataset.

    Parameters
    ----------
    dataset : PupilDataset
        Dataset with augment=True
    n_samples : int
        Number of distinct images to visualize
    n_augs_per_sample : int
        Number of augmented versions per image
    overlay_mask : bool
        If True, overlay mask in red
    """
    fig, axes = plt.subplots(
        n_samples,
        n_augs_per_sample,
        figsize=(3 * n_augs_per_sample, 3 * n_samples),
        squeeze=False,
    )

    sample_indices = np.random.choice(np.arange(len(dataset)), n_samples, replace=False)
    for i in range(n_samples):
        for j in range(n_augs_per_sample):
            img, mask = dataset[sample_indices[i]]
            img_np = img.squeeze().numpy()
            ax = axes[i, j]

            if overlay_mask:
                mask_np = mask.squeeze().numpy()
                rgb = np.stack([img_np] * 3, axis=-1)

                # alpha = 0.35  # mask transparency (0 = invisible, 1 = solid)
                red = np.array([1.0, 0.0, 0.0])

                blended = rgb.copy()
                blended[mask_np > 0] = (1 - mask_transparency) * rgb[
                    mask_np > 0
                ] + mask_transparency * red

                ax.imshow(blended)

            else:
                ax.imshow(img_np, cmap="gray")

            ax.axis("off")

            if i == 0:
                ax.set_title(f"Aug {j+1}", fontsize=10)

    plt.tight_layout()
    plt.show()


# example paths
image_paths = sorted(Path("images_train/").glob("*.png"))
mask_paths = sorted(Path("masks_train/").glob("*.png"))

dataset = PupilDataset(
    image_paths=image_paths,
    mask_paths=mask_paths,
    augment=True,
    target_size=148,
)

show_augmented_samples(
    dataset,
    n_samples=20,
    n_augs_per_sample=2,
    overlay_mask=True,
)