Spaces:
Runtime error
Runtime error
File size: 2,197 Bytes
1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 2cdd41c 1615d09 |
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 |
from pathlib import Path
import cv2
import numpy as np
from isegm.data.base import ISDataset
from isegm.data.sample import DSample
class ImagesDirDataset(ISDataset):
def __init__(
self, dataset_path, images_dir_name="images", masks_dir_name="masks", **kwargs
):
super(ImagesDirDataset, self).__init__(**kwargs)
self.dataset_path = Path(dataset_path)
self._images_path = self.dataset_path / images_dir_name
self._insts_path = self.dataset_path / masks_dir_name
images_list = [x for x in sorted(self._images_path.glob("*.*"))]
samples = {x.stem: {"image": x, "masks": []} for x in images_list}
for mask_path in self._insts_path.glob("*.*"):
mask_name = mask_path.stem
if mask_name in samples:
samples[mask_name]["masks"].append(mask_path)
continue
mask_name_split = mask_name.split("_")
if mask_name_split[-1].isdigit():
mask_name = "_".join(mask_name_split[:-1])
assert mask_name in samples
samples[mask_name]["masks"].append(mask_path)
for x in samples.values():
assert len(x["masks"]) > 0, x["image"]
self.dataset_samples = [v for k, v in sorted(samples.items())]
def get_sample(self, index) -> DSample:
sample = self.dataset_samples[index]
image_path = str(sample["image"])
objects = []
ignored_regions = []
masks = []
for indx, mask_path in enumerate(sample["masks"]):
gt_mask = cv2.imread(str(mask_path))[:, :, 0].astype(np.int32)
instances_mask = np.zeros_like(gt_mask)
instances_mask[gt_mask == 128] = 2
instances_mask[gt_mask > 128] = 1
masks.append(instances_mask)
objects.append((indx, 1))
ignored_regions.append((indx, 2))
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return DSample(
image,
np.stack(masks, axis=2),
objects_ids=objects,
ignore_ids=ignored_regions,
sample_id=index,
)
|