<gh_stars>0
import os
import numpy as np
import scipy.misc as m
from PIL import Image
from torch.utils import data
from mypath import Path
from torchvision import transforms
from dataloaders import custom_transforms as tr

class CropWeedSegmentation(data.Dataset):
    NUM_CLASSES = 3

    def __init__(self, args, root='/home/robot/datasets/structured_cwc/', split="train"):

        self.root = root
        self.split = split
        self.args = args
        self.files = {}
        self.nir = args.nir
        self.images_base = os.path.join(self.root, self.split, 'img')
        self.annotations_base = os.path.join(self.root, self.split, 'lbl')

        self.files[split] = self.recursive_glob(rootdir=self.images_base, suffix='.png')

        self.void_classes = []
        self.valid_classes = [0, 1, 2]
        self.class_names = ['background', 'plants', 'weeds']

        self.ignore_index = 255
        self.class_map = dict(zip(self.valid_classes, range(self.NUM_CLASSES)))

        if not self.files[split]:
            raise Exception("No files for split=[%s] found in %s" % (split, self.images_base))

        print("Found %d %s images" % (len(self.files[split]), split))

    def __len__(self):
        return len(self.files[self.split])

    def __getitem__(self, index):

        img_path = self.files[self.split][index].rstrip()
        lbl_path = os.path.join(self.annotations_base,
                                os.path.basename(img_path)[:-4] + '.png')
        if self.nir:
            _img = Image.open(img_path).convert('RGB')  # width x height x 3
            nir_image = os.path.join('/home/robot/datasets/cwc_dataset/images/nir/',
                                     os.path.basename(img_path)[:-4] + '.png')
            _nir = Image.open(nir_image).convert('L')  # width x height x 1
            _img.putalpha(_nir)  # width x height x 4
        else:
            _img = Image.open(img_path).convert('RGB')  # width x height x 3

        _tmp = np.array(Image.open(lbl_path), dtype=np.uint8)
        _tmp[_tmp == 255] = 1
        _tmp[_tmp == 0] = 0
        _tmp[_tmp == 128] = 2

        _tmp = self.encode_segmap(_tmp)
        _target = Image.fromarray(_tmp)

        sample = {'image': _img, 'label': _target}

        if self.split == 'train':
            return self.transform_tr(sample)
        elif self.split == 'valid':
            return self.transform_val(sample)
        elif self.split == 'test':
            return self.transform_ts(sample)

    def encode_segmap(self, mask):
        # Put all void classes to zero
        for _voidc in self.void_classes:
            mask[mask == _voidc] = self.ignore_index
        for _validc in self.valid_classes:
            mask[mask == _validc] = self.class_map[_validc]
        return mask

    def recursive_glob(self, rootdir='.', suffix=''):
        """Performs recursive glob with given suffix and rootdir
            :param rootdir is the root directory
            :param suffix is the suffix to be searched
        """
        return [os.path.join(looproot, filename)
                for looproot, _, filenames in os.walk(rootdir)
                for filename in filenames if filename.endswith(suffix)]

    def transform_tr(self, sample):
        if self.nir:
            mean = (0.485, 0.456, 0.406, 0.424)         # TODO use another mean[3] value?
            std = (0.229, 0.224, 0.225, 0.221)
        else:
            mean = (0.485, 0.456, 0.406)
            std = (0.229, 0.224, 0.225)
        composed_transforms = transforms.Compose([
            tr.RandomHorizontalFlip(),
            tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size, fill=255),
            tr.RandomGaussianBlur(),
            tr.Normalize(mean=mean, std=std),
            tr.ToTensor()])

        return composed_transforms(sample)

    def transform_val(self, sample):
        if self.nir:
            mean = (0.485, 0.456, 0.406, 0.424)         # TODO use another mean[3] value?
            std = (0.229, 0.224, 0.225, 0.221)
        else:
            mean = (0.485, 0.456, 0.406)
            std = (0.229, 0.224, 0.225)
        composed_transforms = transforms.Compose([
            tr.FixScaleCrop(crop_size=self.args.crop_size),
            tr.Normalize(mean=mean, std=std),
            tr.ToTensor()])

        return composed_transforms(sample)

    def transform_ts(self, sample):
        if self.nir:
            mean = (0.485, 0.456, 0.406, 0.424)         # TODO use another mean[3] value?
            std = (0.229, 0.224, 0.225, 0.221)
        else:
            mean = (0.485, 0.456, 0.406)
            std = (0.229, 0.224, 0.225)
        composed_transforms = transforms.Compose([
            tr.FixedResize(size=self.args.crop_size),
            tr.Normalize(mean=mean, std=std),
            tr.ToTensor()])

        return composed_transforms(sample)

if __name__ == '__main__':
    from dataloaders.utils import decode_segmap
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.base_size = 1296
    # args.base_size = 384

    args.crop_size = 512

    cityscapes_train = CropWeedSegmentation(args, split='train')

    dataloader = DataLoader(cityscapes_train, batch_size=2, shuffle=True, num_workers=2)

    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()
            gt = sample['label'].numpy()
            tmp = np.array(gt[jj]).astype(np.uint8)
            segmap = decode_segmap(tmp, dataset='cityscapes')
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0])
            img_tmp *= (0.229, 0.224, 0.225)
            img_tmp += (0.485, 0.456, 0.406)
            img_tmp *= 255.0
            img_tmp = img_tmp.astype(np.uint8)
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            plt.imshow(segmap)

        if ii == 1:
            break

    plt.show(block=True)


import os
import numpy as np
import scipy.misc as m
from PIL import Image
from torch.utils import data
from mypath import Path
from torchvision import transforms
from dataloaders import custom_transforms as tr

class CropWeedSegmentation(data.Dataset):
    NUM_CLASSES = 3

    def __init__(self, args, root='/home/robot/datasets/structured_cwc/', split="train"):

        self.root = root
        self.split = split
        self.args = args
        self.files = {}
        self.nir = args.nir
        self.images_base = os.path.join(self.root, self.split, 'img')
        self.annotations_base = os.path.join(self.root, self.split, 'lbl')

        self.files[split] = self.recursive_glob(rootdir=self.images_base, suffix='.png')

        self.void_classes = []
        self.valid_classes = [0, 1, 2]
        self.class_names = ['background', 'plants', 'weeds']

        self.ignore_index = 255
        self.class_map = dict(zip(self.valid_classes, range(self.NUM_CLASSES)))

        if not self.files[split]:
            raise Exception("No files for split=[%s] found in %s" % (split, self.images_base))

        print("Found %d %s images" % (len(self.files[split]), split))

    def __len__(self):
        return len(self.files[self.split])

    def __getitem__(self, index):

        img_path = self.files[self.split][index].rstrip()
        lbl_path = os.path.join(self.annotations_base,
                                os.path.basename(img_path)[:-4] + '.png')
        if self.nir:
            _img = Image.open(img_path).convert('RGB')  # width x height x 3
            nir_image = os.path.join('/home/robot/datasets/cwc_dataset/images/nir/',
                                     os.path.basename(img_path)[:-4] + '.png')
            _nir = Image.open(nir_image).convert('L')  # width x height x 1
            _img.putalpha(_nir)  # width x height x 4
        else:
            _img = Image.open(img_path).convert('RGB')  # width x height x 3

        _tmp = np.array(Image.open(lbl_path), dtype=np.uint8)
        _tmp[_tmp == 255] = 1
        _tmp[_tmp == 0] = 0
        _tmp[_tmp == 128] = 2

        _tmp = self.encode_segmap(_tmp)
        _target = Image.fromarray(_tmp)

        sample = {'image': _img, 'label': _target}

        if self.split == 'train':
            return self.transform_tr(sample)
        elif self.split == 'valid':
            return self.transform_val(sample)
        elif self.split == 'test':
            return self.transform_ts(sample)

    def encode_segmap(self, mask):
        # Put all void classes to zero
        for _voidc in self.void_classes:
            mask[mask == _voidc] = self.ignore_index
        for _validc in self.valid_classes:
            mask[mask == _validc] = self.class_map[_validc]
        return mask

    def recursive_glob(self, rootdir='.', suffix=''):
        """Performs recursive glob with given suffix and rootdir
            :param rootdir is the root directory
            :param suffix is the suffix to be searched
        """
        return [os.path.join(looproot, filename)
                for looproot, _, filenames in os.walk(rootdir)
                for filename in filenames if filename.endswith(suffix)]

    def transform_tr(self, sample):
        if self.nir:
            mean = (0.485, 0.456, 0.406, 0.424)         # TODO use another mean[3] value?
            std = (0.229, 0.224, 0.225, 0.221)
        else:
            mean = (0.485, 0.456, 0.406)
            std = (0.229, 0.224, 0.225)
        composed_transforms = transforms.Compose([
            tr.RandomHorizontalFlip(),
            tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size, fill=255),
            tr.RandomGaussianBlur(),
            tr.Normalize(mean=mean, std=std),
            tr.ToTensor()])

        return composed_transforms(sample)

    def transform_val(self, sample):
        if self.nir:
            mean = (0.485, 0.456, 0.406, 0.424)         # TODO use another mean[3] value?
            std = (0.229, 0.224, 0.225, 0.221)
        else:
            mean = (0.485, 0.456, 0.406)
            std = (0.229, 0.224, 0.225)
        composed_transforms = transforms.Compose([
            tr.FixScaleCrop(crop_size=self.args.crop_size),
            tr.Normalize(mean=mean, std=std),
            tr.ToTensor()])

        return composed_transforms(sample)

    def transform_ts(self, sample):
        if self.nir:
            mean = (0.485, 0.456, 0.406, 0.424)         # TODO use another mean[3] value?
            std = (0.229, 0.224, 0.225, 0.221)
        else:
            mean = (0.485, 0.456, 0.406)
            std = (0.229, 0.224, 0.225)
        composed_transforms = transforms.Compose([
            tr.FixedResize(size=self.args.crop_size),
            tr.Normalize(mean=mean, std=std),
            tr.ToTensor()])

        return composed_transforms(sample)

if __name__ == '__main__':
    from dataloaders.utils import decode_segmap
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.base_size = 1296
    # args.base_size = 384

    args.crop_size = 512

    cityscapes_train = CropWeedSegmentation(args, split='train')

    dataloader = DataLoader(cityscapes_train, batch_size=2, shuffle=True, num_workers=2)

    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()
            gt = sample['label'].numpy()
            tmp = np.array(gt[jj]).astype(np.uint8)
            segmap = decode_segmap(tmp, dataset='cityscapes')
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0])
            img_tmp *= (0.229, 0.224, 0.225)
            img_tmp += (0.485, 0.456, 0.406)
            img_tmp *= 255.0
            img_tmp = img_tmp.astype(np.uint8)
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            plt.imshow(segmap)

        if ii == 1:
            break

    plt.show(block=True)
# width x height x 3 # width x height x 1 # width x height x 4 # width x height x 3 # Put all void classes to zero Performs recursive glob with given suffix and rootdir             :param rootdir is the root directory             :param suffix is the suffix to be searched # TODO use another mean[3] value? # TODO use another mean[3] value? # TODO use another mean[3] value? # args.base_size = 384