|
- import os
- from PIL import Image
- import random
- import numpy as np
- import torch
- import torchvision.transforms as tfs
- import random
- import cv2
-
- class columbia_test(object):
- def __init__(self, newsize=320, k=3):
- super(columbia_test, self).__init__()
-
- self.images = []
- self.masks = []
-
- self.newsize = newsize
- self.k = k
- self.noise = gaussian_noise(0, k)
-
- self.im_tfs = tfs.Compose([tfs.ToTensor()])
-
- self.readImage_columbia("/home/gaojintong/all_txt/columbia_test.txt")
- #self.readImage_columbia("/home/gaojintong/all_txt/columbia_180.txt")
-
- print("images number:%d"%(len(self.images)))
-
- def readImage_columbia(self, path):
-
- lines = []
- with open(path, "r") as f:
- for line in f.readlines():
- line = line.strip('\n')
- lines.append(line)
- random.shuffle(lines)
-
- for line in lines:
- self.images.append("/home/gaojintong/dataset/columbia/images/" + line)
- #self.images.append("/home/gaojintong/ttt/jpeg/columbia/" + line[:-4] + '_' + str(self.k) + '.jpg')
- self.masks.append("/home/gaojintong/dataset/columbia/masks/" + line)
-
- def image_transform(self, img, img_gt):
- img = self.im_tfs(img)
- img_gt = self.im_tfs(img_gt)
- return img, img_gt
-
- def __getitem__(self, idx):
- '''
- img = cv2.imread(self.images[idx])
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- img = cv2.resize(img, (self.newsize, self.newsize))
- #img = cv2.GaussianBlur(img,(self.k,self.k),0)
- '''
- img = Image.open(self.images[idx]).convert('RGB').resize((self.newsize, self.newsize))
- img_gt = Image.open(self.masks[idx]).convert('L').resize((self.newsize, self.newsize))
-
- #img = self.noise(img)
-
- img, img_gt = self.image_transform(img, img_gt)
- return img, img_gt
-
- def __len__(self):
- return len(self.images)
-
- class gaussian_noise(object):
-
- def __init__(self, mean, sigma):
-
- self.mean = mean
- self.sigma = sigma
-
- def __call__(self, img):
- """
- args:
- img (pil image): pil image
- returns:
- pil image: pil image.
- """
- img_ = np.array(img).copy()
- img_ = img_ / 255.0
- noise = np.random.normal(self.mean, self.sigma, img_.shape)
- gaussian_out = img_ + noise
- gaussian_out = np.clip(gaussian_out, 0, 1)
- gaussian_out = np.uint8(gaussian_out*255)
- return Image.fromarray(gaussian_out)
-
- if __name__ == '__main__':
-
- name = columbia_test(newsize=320)
- data_loader = torch.utils.data.DataLoader(name, batch_size=10, shuffle=True)
- for data in data_loader:
- img, img_gt = data
- print(img_gt.shape)
-
|