Spaces:
Sleeping
Sleeping
File size: 2,757 Bytes
02ba63a |
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 |
import torch
import torchvision
from torchvision import transforms
import random
from PIL import Image
import os
from utils import RandomAffineAndRetMat
def load_filenames(data_dir):
# label_data = pd.read_json(INPUT_DIR+'DataList.json')
# label_data = label_data.sort_index()
# tmp_points = []
# filenames = []
# for o in tqdm(label_data.data[0:1000]):
# filenames.append(o['filename'])
# a = o['filename']
# tmps = []
# for i in range(60):
# tmps.append(o['points'][str(i)]['x'])
# tmps.append(o['points'][str(i)]['y'])
# tmp_points.append(tmps) # datanum
# filenames = pd.Series(filenames)
# filenames = [str(i).zfill(4)+'.jpg' for i in range(3400)]
# df_points = pd.DataFrame(tmp_points)
# load from data_dir
# 画像の拡張子のみ
img_exts = ['.jpg', '.jpeg', '.png', '.bmp', '.ppm', '.pgm', '.tif', '.tiff']
filenames = [f for f in os.listdir(data_dir) if os.path.splitext(f)[1].lower() in img_exts]
return filenames
class MyDataset:
def __init__(self, X, valid=False, img_dir='resources/trainB/', img_size=256):
self.X = X
self.valid = valid
self.img_dir = img_dir
self.img_size = img_size
def __len__(self):
return len(self.X)
def __getitem__(self, index):
# 画像を読み込んでトランスフォームを適用
f = self.img_dir + self.X[index]
original_X = Image.open(f)
trans = [
transforms.ToTensor(),
# transforms.Normalize(mean=means, std=stds),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.15),
transforms.RandomGrayscale(0.3),
]
transform = transforms.Compose(trans)
xlist = []
matlist = []
is_flip = random.randint(0, 1) # 同じ画像はフリップ
for i in range(2):
af = RandomAffineAndRetMat(
degrees=[-30, 30],
translate=(0.1, 0.1), scale=(0.8, 1.2),
# fill=(random.random(), random.random(), random.random()),
fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
shear=[-10, 10],
interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
)
X, affine_matrix = af(transforms.Resize(self.img_size)(original_X))
# randomflip
if is_flip == 1:
X = transforms.RandomHorizontalFlip(1.)(X)
flip_matrix = torch.tensor([[-1., 0., 0.],
[0., 1., 0.],
[0., 0., 1.]])
affine_matrix = torch.matmul(flip_matrix, affine_matrix)
xlist.append(transform(X))
matlist.append(affine_matrix)
X = torch.stack(xlist)
mat = torch.stack(matlist)
return X, mat, f |