animeinsseg/data/maskrefine_dataset.py

import albumentations as A

from torch.utils.data import Dataset, DataLoader
import pycocotools.mask as maskUtils
from pycocotools.coco import COCO
import random
import os.path as osp
import cv2
import numpy as np
from scipy.ndimage import distance_transform_bf, distance_transform_edt, distance_transform_cdt


def is_grey(img: np.ndarray):
    if len(img.shape) == 3 and img.shape[2] == 3:
        return False
    else:
        return True


def square_pad_resize(img: np.ndarray, tgt_size: int, pad_value = (0, 0, 0)):
    h, w = img.shape[:2]
    pad_h, pad_w = 0, 0
    
    # make square image
    if w < h:
        pad_w = h - w
        w += pad_w
    elif h < w:
        pad_h = w - h
        h += pad_h

    pad_size = tgt_size - h
    if pad_size > 0:
        pad_h += pad_size
        pad_w += pad_size

    if pad_h > 0 or pad_w > 0:    
        c = 1
        if is_grey(img):
            if isinstance(pad_value, tuple):
                pad_value = pad_value[0]
        else:
            if isinstance(pad_value, int):
                pad_value = (pad_value, pad_value, pad_value)

        img = cv2.copyMakeBorder(img, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=pad_value)

    resize_ratio = tgt_size / img.shape[0]
    if resize_ratio < 1:
        img = cv2.resize(img, (tgt_size, tgt_size), interpolation=cv2.INTER_AREA)
    elif resize_ratio > 1:
        img = cv2.resize(img, (tgt_size, tgt_size), interpolation=cv2.INTER_LINEAR)
        
    return img, resize_ratio, pad_h, pad_w


class MaskRefineDataset(Dataset):

    def __init__(self, 
                 refine_ann_path: str, 
                 data_root: str, 
                 load_instance_mask: bool = True, 
                 aug_ins_prob: float = 0.,
                 ins_rect_prob: float = 0.,
                 output_size: int = 720,
                 augmentation: bool = False,
                 with_distance: bool = False):
        self.load_instance_mask = load_instance_mask
        self.ann_util = COCO(refine_ann_path)
        self.img_ids = self.ann_util.getImgIds()
        self.set_load_method(load_instance_mask)
        self.data_root = data_root

        self.ins_rect_prob = ins_rect_prob
        self.aug_ins_prob = aug_ins_prob
        self.augmentation = augmentation
        if augmentation:
            transform = [
                A.OpticalDistortion(),
                A.HorizontalFlip(),
                A.CLAHE(),
                A.Posterize(),
                A.CropAndPad(percent=0.1, p=0.3, pad_mode=cv2.BORDER_CONSTANT, pad_cval=0, pad_cval_mask=0, keep_size=True),
                A.RandomContrast(), 
                A.Rotate(30, p=0.3, mask_value=0, border_mode=cv2.BORDER_CONSTANT)
            ]
            self._aug_transform = A.Compose(transform)
        else:
            self._aug_transform = None

        self.output_size = output_size
        self.with_distance = with_distance

    def set_output_size(self, size: int):
        self.output_size = size

    def set_load_method(self, load_instance_mask: bool):
        if load_instance_mask:
            self._load_mask = self._load_with_instance
        else:
            self._load_mask = self._load_without_instance

    def __getitem__(self, idx: int):
        img_id = self.img_ids[idx]
        img_meta = self.ann_util.imgs[img_id]
        img_path = osp.join(self.data_root, img_meta['file_name'])
        img = cv2.imread(img_path)

        annids = self.ann_util.getAnnIds([img_id])
        if len(annids) > 0:
            ann = random.choice(annids)
            ann = self.ann_util.anns[ann]
            assert ann['image_id'] == img_id
        else:
            ann = None
        
        return self._load_mask(img, ann)
    
    def transform(self, img: np.ndarray, mask: np.ndarray, ins_seg: np.ndarray = None) -> dict:
        if ins_seg is not None:
            use_seg = True
        else:
            use_seg = False

        if self.augmentation:
            masks = [mask]
            if use_seg:
                masks.append(ins_seg)
            data = self._aug_transform(image=img, masks=masks)
            img = data['image']
            masks = data['masks']
            mask = masks[0]
            if use_seg:
                ins_seg = masks[1]

        img = square_pad_resize(img, self.output_size, random.randint(0, 255))[0]
        mask = square_pad_resize(mask, self.output_size, 0)[0]
        if ins_seg is not None:
            ins_seg = square_pad_resize(ins_seg, self.output_size, 0)[0]

        img = (img.astype(np.float32) / 255.).transpose((2, 0, 1))
        mask = mask[None, ...]


        if use_seg:
            ins_seg = ins_seg[None, ...]
            img = np.concatenate((img, ins_seg), axis=0)

        data = {'img': img, 'mask': mask}
        if self.with_distance:
            dist = distance_transform_edt(mask[0])
            dist_max = dist.max()
            if dist_max != 0:
                dist = 1 - dist / dist_max
                # diff_mat = cv2.bitwise_xor(mask[0], ins_seg[0])
                # dist = dist + diff_mat + 0.2
                dist = dist + 0.2
                dist = dist.size / (dist.sum() + 1) * dist
                dist = np.clip(dist, 0, 20)
            else:
                dist = np.ones_like(dist)
                # print(dist.max(), dist.min())
            data['dist_weight'] = dist[None, ...]
        return data

    def _load_with_instance(self, img: np.ndarray, ann: dict):
        if ann is None:
            mask = np.zeros(img.shape[:2], dtype=np.float32)
            ins_seg = mask
        else:
            mask = maskUtils.decode(ann['segmentation']).astype(np.float32)
            if self.augmentation and random.random() < self.ins_rect_prob:
                ins_seg = np.zeros_like(mask)
                bbox = [int(b) for b in ann['bbox']]
                ins_seg[bbox[1]: bbox[1] + bbox[3], bbox[0]: bbox[0] + bbox[2]] = 1
            elif len(ann['pred_segmentations']) > 0:
                ins_seg = random.choice(ann['pred_segmentations'])
                ins_seg = maskUtils.decode(ins_seg).astype(np.float32)
            else:
                ins_seg = mask
            if self.augmentation and random.random() < self.aug_ins_prob:
                ksize = random.choice([1, 3, 5, 7])
                ksize = ksize * 2 + 1
                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(ksize, ksize))
                if random.random() < 0.5:
                    ins_seg = cv2.dilate(ins_seg, kernel)
                else:
                    ins_seg = cv2.erode(ins_seg, kernel)

        return self.transform(img, mask, ins_seg)

    def _load_without_instance(self, img: np.ndarray, ann: dict):
        if ann is None:
            mask = np.zeros(img.shape[:2], dtype=np.float32)
        else:
            mask = maskUtils.decode(ann['segmentation']).astype(np.float32)
        return self.transform(img, mask)

    def __len__(self):
        return len(self.img_ids)


if __name__ == '__main__':
    ann_path = r'workspace/test_syndata/annotations/refine_train.json'
    data_root = r'workspace/test_syndata/train'

    ann_path = r'workspace/test_syndata/annotations/refine_train.json'
    data_root = r'workspace/test_syndata/train'
    aug_ins_prob = 0.5
    load_instance_mask = True
    ins_rect_prob = 0.25
    output_size = 640
    augmentation = True

    random.seed(0)

    md = MaskRefineDataset(ann_path, data_root, load_instance_mask, aug_ins_prob, ins_rect_prob, output_size, augmentation, with_distance=True)
    
    dl = DataLoader(md, batch_size=1, shuffle=False, persistent_workers=True,
                                  num_workers=1, pin_memory=True)
    for data in dl:
        img = data['img'].cpu().numpy()
        img = (img[0, :3].transpose((1, 2, 0)) * 255).astype(np.uint8)
        mask = (data['mask'].cpu().numpy()[0][0] * 255).astype(np.uint8)
        if load_instance_mask:
            ins = (data['img'].cpu().numpy()[0][3] * 255).astype(np.uint8)
            cv2.imshow('ins', ins)
        dist = data['dist_weight'].cpu().numpy()[0][0]
        dist = (dist / dist.max() * 255).astype(np.uint8)
        cv2.imshow('img', img)
        cv2.imshow('mask', mask)
        cv2.imshow('dist_weight', dist)
        cv2.waitKey(0)

        # cv2.imwrite('')