in scripts/data_preparation/prepare_vistas.py [0:0]
def __call__(self, img_id):
coco_ann = []
# Load the annotation
with Image.open(path.join(self.root_dir, _LABELS_DIR, img_id + "." + _LABELS_EXT)) as lbl_img:
lbl = np.array(lbl_img, dtype=np.uint16)
lbl_size = lbl_img.size
mvd_ids = np.unique(lbl)
# Compress the labels and compute cat
lbl_out = np.zeros(lbl.shape, np.int32)
cat = [255]
iscrowd = [0]
for mvd_id in mvd_ids:
mvd_class_id = int(mvd_id // 255)
category = self.categories[mvd_class_id]
# If it's a void class just skip it
if not category["evaluate"]:
continue
# Extract all necessary information
iss_class_id = self.cat_id_mvd_to_iss[mvd_class_id]
iss_instance_id = len(cat)
iscrowd_i = 1 if "group" in category["name"] else 0
mask_i = lbl == mvd_id
# Save ISS format annotation
cat.append(iss_class_id)
iscrowd.append(iscrowd_i)
lbl_out[mask_i] = iss_instance_id
# Compute COCO detection format annotation
if category["instances"]:
category_info = {"id": iss_class_id, "is_crowd": iscrowd_i == 1}
coco_ann_i = pct.create_annotation_info(
counter.increment(), img_id, category_info, mask_i, lbl_size, tolerance=2)
if coco_ann_i is not None:
coco_ann.append(coco_ann_i)
# COCO detection format image annotation
coco_img = pct.create_image_info(img_id, img_id + "." + _IMAGES_EXT, lbl_size)
# Write output
out_msk_dir = path.join(self.out_dir, "msk")
out_img_dir = path.join(self.out_dir, "img")
_ensure_dir(out_msk_dir)
_ensure_dir(out_img_dir)
Image.fromarray(lbl_out).save(path.join(out_msk_dir, img_id + ".png"))
shutil.copy(path.join(self.root_dir, _IMAGES_DIR, img_id + "." + _IMAGES_EXT),
path.join(out_img_dir, img_id + "." + _IMAGES_EXT))
img_meta = {
"id": img_id,
"cat": cat,
"size": (lbl_size[1], lbl_size[0]),
"iscrowd": iscrowd
}
return img_meta, coco_img, coco_ann