|
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 = 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 |
|
|
