Skip to content

dataset_utils

Utility functions for training datasets in the YOLO format.

clean_yolo_dataset(path, img_ext)

Remove images and labels when the label file is empty.

Parameters

path : pathlib.Path The path to the dataset. img_ext : str

Returns

None

Source code in chirpdetector/dataset_utils.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def clean_yolo_dataset(path: pathlib.Path, img_ext: str) -> None:
    """Remove images and labels when the label file is empty.

    Parameters
    ----------
    path : pathlib.Path
        The path to the dataset.
    img_ext : str

    Returns
    -------
    None
    """
    img_path = path / "images"
    lbl_path = path / "labels"

    images = list(img_path.glob(f"*{img_ext}"))

    for image in images:
        lbl = lbl_path / f"{image.stem}.txt"
        if lbl.stat().st_size == 0:
            image.unlink()
            lbl.unlink()

load_img(path)

Load an image from a path as a numpy array.

Parameters

path : pathlib.Path The path to the image.

Returns

img : np.ndarray The image as a numpy array.

Source code in chirpdetector/dataset_utils.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def load_img(path: pathlib.Path) -> np.ndarray:
    """Load an image from a path as a numpy array.

    Parameters
    ----------
    path : pathlib.Path
        The path to the image.

    Returns
    -------
    img : np.ndarray
        The image as a numpy array.
    """
    img = Image.open(path)
    return np.asarray(img)

merge_yolo_datasets(dataset1, dataset2, output)

Merge two yolo-style datasets into one.

Parameters

dataset1 : str The path to the first dataset. dataset2 : str The path to the second dataset. output : str The path to the output dataset.

Returns

None

Source code in chirpdetector/dataset_utils.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def merge_yolo_datasets(
    dataset1: pathlib.Path,
    dataset2: pathlib.Path,
    output: pathlib.Path,
) -> None:
    """Merge two yolo-style datasets into one.

    Parameters
    ----------
    dataset1 : str
        The path to the first dataset.
    dataset2 : str
        The path to the second dataset.
    output : str
        The path to the output dataset.

    Returns
    -------
    None
    """
    dataset1 = pathlib.Path(dataset1)
    dataset2 = pathlib.Path(dataset2)
    output = pathlib.Path(output)

    if not dataset1.exists():
        msg = f"{dataset1} does not exist."
        raise FileNotFoundError(msg)
    if not dataset2.exists():
        msg = f"{dataset2} does not exist."
        raise FileNotFoundError(msg)
    if output.exists():
        msg = f"{output} already exists."
        raise FileExistsError(msg)

    output_images = output / "images"
    output_images.mkdir(parents=True, exist_ok=False)
    output_labels = output / "labels"
    output_labels.mkdir(parents=True, exist_ok=False)

    imgs1 = list((dataset1 / "images").iterdir())
    labels1 = list((dataset1 / "labels").iterdir())
    imgs2 = list((dataset2 / "images").iterdir())
    labels2 = list((dataset2 / "labels").iterdir())

    print(f"Found {len(imgs1)} images in {dataset1}.")
    print(f"Found {len(imgs2)} images in {dataset2}.")

    print(f"Copying images and labels to {output}...")
    for idx, _ in enumerate(imgs1):
        shutil.copy(imgs1[idx], output_images / imgs1[idx].name)
        shutil.copy(labels1[idx], output_labels / labels1[idx].name)

    for idx, _ in enumerate(imgs2):
        shutil.copy(imgs2[idx], output_images / imgs2[idx].name)
        shutil.copy(labels2[idx], output_labels / labels2[idx].name)

    classes = dataset1 / "classes.txt"
    shutil.copy(classes, output / classes.name)

    print(f"Done. Merged {len(imgs1) + len(imgs2)} images.")

plot_yolo_dataset(path, n)

Plot n random images YOLO-style dataset.

Parameters

path : pathlib.Path The path to the dataset.

Returns

None

Source code in chirpdetector/dataset_utils.py 
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
def plot_yolo_dataset(path: pathlib.Path, n: int) -> None:
    """Plot n random images YOLO-style dataset.

    Parameters
    ----------
    path : pathlib.Path
        The path to the dataset.

    Returns
    -------
    None
    """
    mpl.use("TkAgg")
    labelpath = path / "labels"
    imgpath = path / "images"

    label_paths = np.array(list(labelpath.glob("*.txt")))
    rng = np.random.default_rng()
    label_paths = rng.choice(label_paths, size=n, replace=False)

    for lp in label_paths:
        imgp = imgpath / (lp.stem + ".png")
        img = load_img(imgp)
        labs = np.loadtxt(lp, dtype=np.float32).reshape(-1, 5)

        coords = labs[:, 1:]

        # make coords absolute and normalize
        coords[:, 0] *= img.shape[1]
        coords[:, 1] *= img.shape[0]
        coords[:, 2] *= img.shape[1]
        coords[:, 3] *= img.shape[0]

        # turn centerx, centery, width, height into xmin, ymin, xmax, ymax
        xmin = coords[:, 0] - coords[:, 2] / 2
        ymin = coords[:, 1] - coords[:, 3] / 2
        xmax = coords[:, 0] + coords[:, 2] / 2
        ymax = coords[:, 1] + coords[:, 3] / 2

        # plot the image
        _, ax = plt.subplots(figsize=(15, 5), constrained_layout=True)
        ax.imshow(img, cmap="magma")
        for i in range(len(xmin)):
            ax.add_patch(
                Rectangle(
                    (xmin[i], ymin[i]),
                    xmax[i] - xmin[i],
                    ymax[i] - ymin[i],
                    fill=False,
                    color="white",
                ),
            )
        ax.set_title(imgp.stem)
        plt.axis("off")
        plt.show()

subset_yolo_dataset(path, img_ext, n)

Subset a YOLO dataset.

Parameters

path : pathlib.Path The path to the dataset root. img_ext : str The image extension, e.g. .png or .jpg n : int The size of the subset

Returns

None

Source code in chirpdetector/dataset_utils.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def subset_yolo_dataset(path: pathlib.Path, img_ext: str, n: int) -> None:
    """Subset a YOLO dataset.

    Parameters
    ----------
    path : pathlib.Path
        The path to the dataset root.
    img_ext : str
        The image extension, e.g. .png or .jpg
    n : int
        The size of the subset

    Returns
    -------
    None
    """
    img_path = path / "images"
    lbl_path = path / "labels"

    images = np.array(img_path.glob(f"*{img_ext}"))
    rng = np.random.default_rng()
    images = rng.choice(images, size=len(images), replace=False)

    images = images[:n]

    subset_dir = path.parent / f"{path.name}_subset"
    subset_dir.mkdir(exist_ok=True)

    subset_img_path = subset_dir / "images"
    subset_img_path.mkdir(exist_ok=True)
    subset_lbl_path = subset_dir / "labels"
    subset_lbl_path.mkdir(exist_ok=True)

    shutil.copy(path / "classes.txt", subset_dir)

    for image in images:
        shutil.copy(image, subset_img_path)
        shutil.copy(lbl_path / f"{image.stem}.txt", subset_lbl_path)