Spaces:
Runtime error
Runtime error
File size: 3,213 Bytes
023485e |
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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from scipy.io import loadmat
from PIL import Image
import numpy as np
import os
from glob import glob
import cv2
dir_name = os.path.dirname(os.path.abspath(__file__))
def cal_new_size(im_h, im_w, min_size, max_size):
if im_h < im_w:
if im_h < min_size:
ratio = 1.0 * min_size / im_h
im_h = min_size
im_w = round(im_w * ratio)
elif im_h > max_size:
ratio = 1.0 * max_size / im_h
im_h = max_size
im_w = round(im_w * ratio)
else:
ratio = 1.0
else:
if im_w < min_size:
ratio = 1.0 * min_size / im_w
im_w = min_size
im_h = round(im_h * ratio)
elif im_w > max_size:
ratio = 1.0 * max_size / im_w
im_w = max_size
im_h = round(im_h * ratio)
else:
ratio = 1.0
return im_h, im_w, ratio
def generate_data(im_path, min_size, max_size):
im = Image.open(im_path)
im_w, im_h = im.size
mat_path = im_path.replace('.jpg', '_ann.mat')
points = loadmat(mat_path)['annPoints'].astype(np.float32)
idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
points = points[idx_mask]
im_h, im_w, rr = cal_new_size(im_h, im_w, min_size, max_size)
im = np.array(im)
if rr != 1.0:
im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
points = points * rr
return Image.fromarray(im), points
def main(input_dataset_path, output_dataset_path, min_size=512, max_size=2048):
for phase in ['Train', 'Test']:
sub_dir = os.path.join(input_dataset_path, phase)
if phase == 'Train':
sub_phase_list = ['train', 'val']
for sub_phase in sub_phase_list:
sub_save_dir = os.path.join(output_dataset_path, sub_phase)
if not os.path.exists(sub_save_dir):
os.makedirs(sub_save_dir)
with open(os.path.join(dir_name, 'qnrf_{}.txt'.format(sub_phase))) as f:
for i in f:
im_path = os.path.join(sub_dir, i.strip())
name = os.path.basename(im_path)
print(name)
im, points = generate_data(im_path, min_size, max_size)
im_save_path = os.path.join(sub_save_dir, name)
im.save(im_save_path)
gd_save_path = im_save_path.replace('jpg', 'npy')
np.save(gd_save_path, points)
else:
sub_save_dir = os.path.join(output_dataset_path, 'test')
if not os.path.exists(sub_save_dir):
os.makedirs(sub_save_dir)
im_list = glob(os.path.join(sub_dir, '*jpg'))
for im_path in im_list:
name = os.path.basename(im_path)
print(name)
im, points = generate_data(im_path, min_size, max_size)
im_save_path = os.path.join(sub_save_dir, name)
im.save(im_save_path)
gd_save_path = im_save_path.replace('jpg', 'npy')
np.save(gd_save_path, points)
|