aliabd/Anime2Sketch

1import os 
2from PIL import Image
3import torchvision.transforms as transforms
4import numpy as np
5import torch 
6
7IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP']
8
9def is_image_file(filename):
10    """if a given filename is a valid image
11    Parameters:
12        filename (str) -- image filename
13    """
14    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
15
16def get_image_list(path):
17    """read the paths of valid images from the given directory path
18    Parameters:
19        path (str)    -- input directory path
20    """
21    assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
22    images = []
23    for dirpath, _, fnames in sorted(os.walk(path)):
24        for fname in sorted(fnames):
25            if is_image_file(fname):
26                img_path = os.path.join(dirpath, fname)
27                images.append(img_path)
28    assert images, '{:s} has no valid image file'.format(path)
29    return images
30
31def get_transform(load_size=0, grayscale=False, method=Image.BICUBIC, convert=True):
32    transform_list = []
33    if grayscale:
34        transform_list.append(transforms.Grayscale(1))
35    if load_size > 0:
36        osize = [load_size, load_size]
37        transform_list.append(transforms.Resize(osize, method))
38    if convert:
39        transform_list += [transforms.ToTensor()]
40        if grayscale:
41            transform_list += [transforms.Normalize((0.5,), (0.5,))]
42        else:
43            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
44    return transforms.Compose(transform_list)
45
46def read_img_path(path, load_size):
47    """read tensors from a given image path
48    Parameters:
49        path (str)     -- input image path
50        load_size(int) -- the input size. If <= 0, don't resize
51    """
52    img = Image.open(path).convert('RGB')
53    aus_resize = None
54    if load_size > 0:
55        aus_resize = img.size
56    transform = get_transform(load_size=load_size)
57    image = transform(img)
58    return image.unsqueeze(0), aus_resize 
59
60def tensor_to_img(input_image, imtype=np.uint8):
61    """"Converts a Tensor array into a numpy image array.
62    Parameters:
63        input_image (tensor) --  the input image tensor array
64        imtype (type)        --  the desired type of the converted numpy array
65    """
66
67    if not isinstance(input_image, np.ndarray):
68        if isinstance(input_image, torch.Tensor):  # get the data from a variable
69            image_tensor = input_image.data
70        else:
71            return input_image
72        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
73        if image_numpy.shape[0] == 1:  # grayscale to RGB
74            image_numpy = np.tile(image_numpy, (3, 1, 1))
75        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
76    else:  # if it is a numpy array, do nothing
77        image_numpy = input_image
78    return image_numpy.astype(imtype)
79
80def save_image(image_numpy, image_path, output_resize=None):
81    """Save a numpy image to the disk
82    Parameters:
83        image_numpy (numpy array)    -- input numpy array
84        image_path (str)             -- the path of the image
85        output_resize(None or tuple) -- the output size. If None, don't resize
86    """
87
88    image_pil = Image.fromarray(image_numpy)
89    if output_resize:
90        image_pil = image_pil.resize(output_resize, Image.BICUBIC)
91    image_pil.save(image_path)