forked from sebasvega95/neural-style-transfer
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
31 lines (23 loc) · 917 Bytes
/
utils.py
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
import cv2
import itertools
import numpy as np
import pathlib
def preprocess_image_from_path(path, target_shape=None):
image = cv2.imread(str(path))
if target_shape is not None:
image = cv2.resize(image, target_shape)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = image.astype(np.float32)
image = np.expand_dims(image, axis=0)
return image
def image_generator(images_path, batch_size=1, target_shape=None):
paths = itertools.cycle(pathlib.Path(images_path).glob('*.jpg'))
while True:
batch_paths = itertools.islice(paths, batch_size)
images = [preprocess_image_from_path(str(path), target_shape=target_shape) for path in batch_paths]
batch = np.concatenate(images, axis=0)
yield batch
def get_num_images(path):
return len(list(pathlib.Path(path).glob('*.jpg')))
def path_exists(path):
return pathlib.Path(path).exists()