|
- import os
- from PIL import Image
- import random
- import numpy as np
- import torch
- import torchvision.transforms as tfs
- import random
- import cv2
-
- class coverage_test(object):
- def __init__(self, newsize=320, k=3):
- super(coverage_test, self).__init__()
-
- self.images = []
- self.masks = []
-
- self.newsize = newsize
- self.k = k
-
- self.im_tfs = tfs.Compose([tfs.ToTensor()])
-
- self.readImage_coverage("/home/gaojintong/all_txt/coverage_test.txt")
- #self.readImage_coverage("/home/gaojintong/all_txt/coverage_100.txt")
-
- print("images number:%d"%(len(self.images)))
-
- def readImage_coverage(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/coverage/images/" + line)
- self.masks.append("/home/gaojintong/dataset/coverage/masks/" + line[:-5] + 'forged.tif')
-
- 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, img_gt = self.image_transform(img, img_gt)
- return img, img_gt
-
- def __len__(self):
- return len(self.images)
-
-
- if __name__ == '__main__':
-
- name = coverage_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)
-
|