mirror of https://github.com/XingangPan/DragGAN
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
38 lines
953 B
Python
38 lines
953 B
Python
# Copyright (c) SenseTime Research. All rights reserved.
|
|
|
|
|
|
import os
|
|
|
|
from PIL import Image
|
|
|
|
IMG_EXTENSIONS = [
|
|
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
|
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
|
|
]
|
|
|
|
|
|
def is_image_file(filename):
|
|
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
|
|
|
|
|
def tensor2im(var):
|
|
# var shape: (3, H, W)
|
|
var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy()
|
|
var = ((var + 1) / 2)
|
|
var[var < 0] = 0
|
|
var[var > 1] = 1
|
|
var = var * 255
|
|
return Image.fromarray(var.astype('uint8'))
|
|
|
|
|
|
def make_dataset(dir):
|
|
images = []
|
|
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
|
for root, _, fnames in sorted(os.walk(dir)):
|
|
for fname in fnames:
|
|
if is_image_file(fname):
|
|
path = os.path.join(root, fname)
|
|
fname = fname.split('.')[0]
|
|
images.append((fname, path))
|
|
return images
|