REVA-QCAV/utils/load.py

47 lines
1.3 KiB
Python
Raw Normal View History

#
# load.py : utils on generators / lists of ids to transform from strings to
# cropped images and masks
2017-08-17 13:32:38 +00:00
import os
import numpy as np
2017-08-17 13:32:38 +00:00
from PIL import Image
from .utils import resize_and_crop, get_square, normalize, hwc_to_chw
2017-08-17 13:32:38 +00:00
def get_ids(dir):
"""Returns a list of the ids in the directory"""
return (f[:-4] for f in os.listdir(dir))
2017-08-17 13:32:38 +00:00
def split_ids(ids, n=2):
"""Split each id in n, creating n tuples (id, k) for each id"""
return ((id, i) for id in ids for i in range(n))
2017-08-17 13:32:38 +00:00
def to_cropped_imgs(ids, dir, suffix, scale):
"""From a list of tuples, returns the correct cropped img"""
2017-08-17 13:32:38 +00:00
for id, pos in ids:
im = resize_and_crop(Image.open(dir + id + suffix), scale=scale)
2017-08-17 13:32:38 +00:00
yield get_square(im, pos)
def get_imgs_and_masks(ids, dir_img, dir_mask, scale):
"""Return all the couples (img, mask)"""
2017-08-17 13:32:38 +00:00
imgs = to_cropped_imgs(ids, dir_img, '.jpg', scale)
2017-08-17 13:32:38 +00:00
# need to transform from HWC to CHW
imgs_switched = map(hwc_to_chw, imgs)
imgs_normalized = map(normalize, imgs_switched)
2017-08-17 13:32:38 +00:00
masks = to_cropped_imgs(ids, dir_mask, '_mask.gif', scale)
2017-08-17 13:32:38 +00:00
return zip(imgs_normalized, masks)
2017-08-21 16:00:07 +00:00
2017-08-21 16:00:07 +00:00
def get_full_img_and_mask(id, dir_img, dir_mask):
im = Image.open(dir_img + id + '.jpg')
mask = Image.open(dir_mask + id + '_mask.gif')
return np.array(im), np.array(mask)