torch_em.util.test

@private

 1"""@private
 2"""
 3import os
 4import imageio
 5import h5py
 6import numpy as np
 7import torch
 8
 9import bioimage_cpp as bic
10
11
12def make_gt(spatial_shape, n_batches=None, with_channels=False, with_background=False, dtype=None):
13    def _make_gt():
14        seeds = np.random.rand(*spatial_shape)
15        seeds = bic.segmentation.label(seeds > 0.99)
16        hmap = bic.distance.distance_transform(seeds == 0)
17        if with_background:
18            mask = np.random.rand(*spatial_shape) > 0.5
19            assert mask.shape == hmap.shape
20        else:
21            mask = None
22        # bic returns uint64 labels; cast to uint32 which torch.from_numpy supports.
23        return bic.segmentation.watershed(hmap, markers=seeds, mask=mask).astype("uint32")
24
25    if n_batches is None and not with_channels:
26        seg = _make_gt()
27    elif n_batches is None and with_channels:
28        seg = _make_gt[None]
29    else:
30        seg = []
31        for _ in range(n_batches):
32            batch_seg = _make_gt()
33            if with_channels:
34                batch_seg = batch_seg[None]
35            seg.append(batch_seg[None])
36        seg = np.concatenate(seg, axis=0)
37    if dtype is not None:
38        seg = seg.astype(dtype)
39    return torch.from_numpy(seg)
40
41
42def create_segmentation_test_data(data_path, raw_key, label_key, shape, chunks):
43    with h5py.File(data_path, "a") as f:
44        f.create_dataset(raw_key, data=np.random.rand(*shape), chunks=chunks)
45        f.create_dataset(label_key, data=np.random.randint(0, 4, size=shape), chunks=chunks)
46
47
48def create_image_collection_test_data(folder, n_images, min_shape, max_shape):
49    im_folder = os.path.join(folder, "images")
50    label_folder = os.path.join(folder, "labels")
51    os.makedirs(im_folder, exist_ok=True)
52    os.makedirs(label_folder, exist_ok=True)
53
54    for i in range(n_images):
55        shape = tuple(np.random.randint(mins, maxs) for mins, maxs in zip(min_shape, max_shape))
56        raw = np.random.rand(*shape).astype("int16")
57        label = np.random.randint(0, 4, size=shape)
58        imageio.imwrite(os.path.join(im_folder, f"im_{i}.tif"), raw)
59        imageio.imwrite(os.path.join(label_folder, f"im_{i}.tif"), label)
def make_gt( spatial_shape, n_batches=None, with_channels=False, with_background=False, dtype=None): 
13def make_gt(spatial_shape, n_batches=None, with_channels=False, with_background=False, dtype=None):
14    def _make_gt():
15        seeds = np.random.rand(*spatial_shape)
16        seeds = bic.segmentation.label(seeds > 0.99)
17        hmap = bic.distance.distance_transform(seeds == 0)
18        if with_background:
19            mask = np.random.rand(*spatial_shape) > 0.5
20            assert mask.shape == hmap.shape
21        else:
22            mask = None
23        # bic returns uint64 labels; cast to uint32 which torch.from_numpy supports.
24        return bic.segmentation.watershed(hmap, markers=seeds, mask=mask).astype("uint32")
25
26    if n_batches is None and not with_channels:
27        seg = _make_gt()
28    elif n_batches is None and with_channels:
29        seg = _make_gt[None]
30    else:
31        seg = []
32        for _ in range(n_batches):
33            batch_seg = _make_gt()
34            if with_channels:
35                batch_seg = batch_seg[None]
36            seg.append(batch_seg[None])
37        seg = np.concatenate(seg, axis=0)
38    if dtype is not None:
39        seg = seg.astype(dtype)
40    return torch.from_numpy(seg)
def create_segmentation_test_data(data_path, raw_key, label_key, shape, chunks): 
43def create_segmentation_test_data(data_path, raw_key, label_key, shape, chunks):
44    with h5py.File(data_path, "a") as f:
45        f.create_dataset(raw_key, data=np.random.rand(*shape), chunks=chunks)
46        f.create_dataset(label_key, data=np.random.randint(0, 4, size=shape), chunks=chunks)
def create_image_collection_test_data(folder, n_images, min_shape, max_shape): 
49def create_image_collection_test_data(folder, n_images, min_shape, max_shape):
50    im_folder = os.path.join(folder, "images")
51    label_folder = os.path.join(folder, "labels")
52    os.makedirs(im_folder, exist_ok=True)
53    os.makedirs(label_folder, exist_ok=True)
54
55    for i in range(n_images):
56        shape = tuple(np.random.randint(mins, maxs) for mins, maxs in zip(min_shape, max_shape))
57        raw = np.random.rand(*shape).astype("int16")
58        label = np.random.randint(0, 4, size=shape)
59        imageio.imwrite(os.path.join(im_folder, f"im_{i}.tif"), raw)
60        imageio.imwrite(os.path.join(label_folder, f"im_{i}.tif"), label)