|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.utils.data as Data |
|
|
import torchvision.transforms as transforms |
|
|
|
|
|
import os |
|
|
from PIL import Image, ImageOps, ImageFilter |
|
|
import os.path as osp |
|
|
import sys |
|
|
import random |
|
|
import shutil |
|
|
|
|
|
|
|
|
class IRSTD_Dataset(Data.Dataset): |
|
|
def __init__(self, args, mode='train'): |
|
|
|
|
|
dataset_dir = args.dataset_dir |
|
|
|
|
|
if mode == 'train': |
|
|
txtfile = 'trainval.txt' |
|
|
elif mode == 'val': |
|
|
txtfile = 'test.txt' |
|
|
|
|
|
self.list_dir = osp.join(dataset_dir, txtfile) |
|
|
self.imgs_dir = osp.join(dataset_dir, 'images') |
|
|
self.label_dir = osp.join(dataset_dir, 'masks') |
|
|
|
|
|
self.names = [] |
|
|
with open(self.list_dir, 'r') as f: |
|
|
self.names += [line.strip() for line in f.readlines()] |
|
|
|
|
|
self.mode = mode |
|
|
self.crop_size = args.crop_size |
|
|
self.base_size = args.base_size |
|
|
self.transform = transforms.Compose([ |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([.485, .456, .406], [.229, .224, .225]), |
|
|
]) |
|
|
|
|
|
def __getitem__(self, i): |
|
|
name = self.names[i] |
|
|
img_path = osp.join(self.imgs_dir, name + '.png') |
|
|
label_path = osp.join(self.label_dir, name + '.png') |
|
|
|
|
|
img = Image.open(img_path).convert('RGB') |
|
|
mask = Image.open(label_path) |
|
|
|
|
|
if self.mode == 'train': |
|
|
img, mask = self._sync_transform(img, mask) |
|
|
elif self.mode == 'val': |
|
|
img, mask = self._testval_sync_transform(img, mask) |
|
|
else: |
|
|
raise ValueError("Unkown self.mode") |
|
|
|
|
|
img, mask = self.transform(img), transforms.ToTensor()(mask) |
|
|
return img, mask |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.names) |
|
|
|
|
|
def _sync_transform(self, img, mask): |
|
|
|
|
|
if random.random() < 0.5: |
|
|
img = img.transpose(Image.FLIP_LEFT_RIGHT) |
|
|
mask = mask.transpose(Image.FLIP_LEFT_RIGHT) |
|
|
crop_size = self.crop_size |
|
|
|
|
|
long_size = random.randint(int(self.base_size * 0.5), int(self.base_size * 2.0)) |
|
|
w, h = img.size |
|
|
if h > w: |
|
|
oh = long_size |
|
|
ow = int(1.0 * w * long_size / h + 0.5) |
|
|
short_size = ow |
|
|
else: |
|
|
ow = long_size |
|
|
oh = int(1.0 * h * long_size / w + 0.5) |
|
|
short_size = oh |
|
|
img = img.resize((ow, oh), Image.BILINEAR) |
|
|
mask = mask.resize((ow, oh), Image.NEAREST) |
|
|
|
|
|
if short_size < crop_size: |
|
|
padh = crop_size - oh if oh < crop_size else 0 |
|
|
padw = crop_size - ow if ow < crop_size else 0 |
|
|
img = ImageOps.expand(img, border=(0, 0, padw, padh), fill=0) |
|
|
mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=0) |
|
|
|
|
|
w, h = img.size |
|
|
x1 = random.randint(0, w - crop_size) |
|
|
y1 = random.randint(0, h - crop_size) |
|
|
img = img.crop((x1, y1, x1 + crop_size, y1 + crop_size)) |
|
|
mask = mask.crop((x1, y1, x1 + crop_size, y1 + crop_size)) |
|
|
|
|
|
if random.random() < 0.5: |
|
|
img = img.filter(ImageFilter.GaussianBlur( |
|
|
radius=random.random())) |
|
|
return img, mask |
|
|
|
|
|
def _testval_sync_transform(self, img, mask): |
|
|
base_size = self.base_size |
|
|
img = img.resize((base_size, base_size), Image.BILINEAR) |
|
|
mask = mask.resize((base_size, base_size), Image.NEAREST) |
|
|
|
|
|
return img, mask |
