File size: 3,992 Bytes
29858c0 |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import torch
import numpy as np
from PIL import Image
import numpy.ma as ma
import torch.utils.data as data
import copy
from torchvision import transforms
import scipy.io as scio
import torchvision.datasets as dset
import random
import scipy.misc
import scipy.io as scio
import os
from PIL import ImageEnhance
from PIL import ImageFilter
class SegDataset(data.Dataset):
def __init__(self, root_dir, txtlist, use_noise, length):
self.path = []
self.real_path = []
self.use_noise = use_noise
self.root = root_dir
input_file = open(txtlist)
while 1:
input_line = input_file.readline()
if not input_line:
break
if input_line[-1:] == '\n':
input_line = input_line[:-1]
self.path.append(copy.deepcopy(input_line))
if input_line[:5] == 'data/':
self.real_path.append(copy.deepcopy(input_line))
input_file.close()
self.length = length
self.data_len = len(self.path)
self.back_len = len(self.real_path)
self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
self.back_front = np.array([[1 for i in range(640)] for j in range(480)])
def __getitem__(self, idx):
index = random.randint(0, self.data_len - 10)
label = np.array(Image.open('{0}/{1}-label.png'.format(self.root, self.path[index])))
meta = scio.loadmat('{0}/{1}-meta.mat'.format(self.root, self.path[index]))
if not self.use_noise:
rgb = np.array(Image.open('{0}/{1}-color.png'.format(self.root, self.path[index])).convert("RGB"))
else:
rgb = np.array(self.trancolor(Image.open('{0}/{1}-color.png'.format(self.root, self.path[index])).convert("RGB")))
if self.path[index][:8] == 'data_syn':
rgb = Image.open('{0}/{1}-color.png'.format(self.root, self.path[index])).convert("RGB")
rgb = ImageEnhance.Brightness(rgb).enhance(1.5).filter(ImageFilter.GaussianBlur(radius=0.8))
rgb = np.array(self.trancolor(rgb))
seed = random.randint(0, self.back_len - 10)
back = np.array(self.trancolor(Image.open('{0}/{1}-color.png'.format(self.root, self.path[seed])).convert("RGB")))
back_label = np.array(Image.open('{0}/{1}-label.png'.format(self.root, self.path[seed])))
mask = ma.getmaskarray(ma.masked_equal(label, 0))
back = np.transpose(back, (2, 0, 1))
rgb = np.transpose(rgb, (2, 0, 1))
rgb = rgb + np.random.normal(loc=0.0, scale=5.0, size=rgb.shape)
rgb = back * mask + rgb
label = back_label * mask + label
rgb = np.transpose(rgb, (1, 2, 0))
#scipy.misc.imsave('embedding_final/rgb_{0}.png'.format(index), rgb)
#scipy.misc.imsave('embedding_final/label_{0}.png'.format(index), label)
if self.use_noise:
choice = random.randint(0, 3)
if choice == 0:
rgb = np.fliplr(rgb)
label = np.fliplr(label)
elif choice == 1:
rgb = np.flipud(rgb)
label = np.flipud(label)
elif choice == 2:
rgb = np.fliplr(rgb)
rgb = np.flipud(rgb)
label = np.fliplr(label)
label = np.flipud(label)
obj = meta['cls_indexes'].flatten().astype(np.int32)
obj = np.append(obj, [0], axis=0)
target = copy.deepcopy(label)
rgb = np.transpose(rgb, (2, 0, 1))
rgb = self.norm(torch.from_numpy(rgb.astype(np.float32)))
target = torch.from_numpy(target.astype(np.int64))
return rgb, target
def __len__(self):
return self.length
|