def simplify(img, hed_model_path):
w, h = img.width, img.height
size_thresh = 0.001 * w * h
img = pil2cv(img)
img = cv2.GaussianBlur(img, (3, 3), 0)
img = cv2.GaussianBlur(img, (3, 3), 0)
img = hed_processing.run_hed(cv2pil(img), hed_model_path)
ret, img = cv2.threshold(pil2cv(img), 50, 255, 0)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = remove_small_objects(img.astype('bool'), size_thresh)
img = 255 * skeletonize(img).astype('uint8')
img = cv2pil(img)
return img
def main(args):
input_src, shuffle, max_num_images, min_w, min_h, max_w, max_h = args.input_src, args.shuffle, args.max_num_images, args.min_dim, args.min_dim, args.max_dim, args.max_dim
output_dir, out_w, out_h, pct_test, save_mode, save_ext = args.output_dir, args.w, args.h, args.pct_test, args.save_mode, args.save_ext
num_per, frac, frac_vary, max_ang_rot, max_stretch, centered = args.num_per, args.frac, args.frac_vary, args.max_ang_rot, args.max_stretch, args.centered
action, target_face_image, face_crop, face_crop_lerp, landmarks_path, hed_model_path = args.action, args.target_face_image, args.face_crop, args.face_crop_lerp, args.landmarks_path, args.hed_model_path
#os.system('rm -rf %s'%output_dir)
# get list of actions
actions = action.split(',')
if False in [a in allowable_actions for a in actions]:
raise Exception('one of your actions does not exist')
# initialize face_processing if needed
if 'face' in actions:
initialize_face_processing(landmarks_path)
target_encodings = get_encodings(
target_face_image) if target_face_image else None
# initialize photosketch if needed
if 'sketch' in actions:
photosketch_processing.setup(args.photosketch_model_path)
# initialize esrgan if needed
if 'upsample' in actions:
esrgan_processing.setup(args.esrgan_model_path)
# initialize SSS if needed
if 'sss' in actions:
sss_processing.setup(args.sss_model_path)
# setup output directories
if output_dir != 'None':
trainA_dir, trainB_dir, testA_dir, testB_dir = setup_output_dirs(
output_dir, save_mode, pct_test > 0)
# initialize input
ext = os.path.splitext(input_src)[1]
is_movie = ext.lower() in ['.mp4', '.mov', '.avi']
if is_movie:
cap = cv2.VideoCapture(input_src)
fps = video.FPS().start()
num_images = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
pct_frames = list(np.linspace(0, 1, num_images))
all_frames = get_frame_indexes(max_num_images, num_images, shuffle)
else:
images = sorted([
f for f in os.listdir(input_src)
if os.path.isfile(os.path.join(input_src, f))
])
num_images = len(images)
all_frames = get_frame_indexes(max_num_images, num_images, shuffle)
# training/test split
training = [1] * len(all_frames) * num_per
if pct_test > 0:
n_test = int(len(all_frames) * num_per * pct_test)
test_per = 1.0 / pct_test
test_idx = [int(test_per * (i + 1) - 1) for i in range(n_test)]
for t in test_idx:
training[t] = 0
# iterate through each input
print("Iterating through %d input images" % len(all_frames))
for k, idx_frame in tqdm(enumerate(all_frames)):
if is_movie:
pct_frame = pct_frames[idx_frame]
frame = int(pct_frame * num_images)
cap.set(1, frame)
ret, img = cap.read()
frame_name = 'frame%06d' % frame
img = cv2pil(img)
else:
img_path = images[idx_frame]
frame_name = os.path.splitext(img_path)[0]
full_image_path = os.path.join(input_src, img_path)
img = Image.open(full_image_path).convert("RGB")
# skip images which are too small or too big
if img.width < min_w or img.height < min_h:
continue
if img.width > max_w or img.height > max_h:
continue
# first crop around face if requested
if face_crop is not None:
jx, jy, jw, jh = get_crop_around_face(img, target_encodings,
out_w / out_h, face_crop,
face_crop_lerp)
img = img.crop((jx, jy, jx + jw, jy + jh))
# preprocess/augment and produce input images
imgs0, imgs1 = augmentation(img, num_per, out_w, out_h, frac,
frac_vary, max_ang_rot, max_stretch,
centered), []
# process each input image to make output
for img0 in imgs0:
img = img0
for a in actions:
if a == 'segment':
img = segment(img)
elif a == 'colorize':
colors = [[255, 255, 255], [0, 0, 0], [127, 0, 0],
[0, 0, 127], [0, 127, 0]]
img = quantize_colors(img, colors)
elif a == 'trace':
img = trace(img)
elif a == 'hed':
img = hed_processing.run_hed(img, hed_model_path)
elif a == 'sketch':
img = photosketch_processing.sketch(img)
elif a == 'simplify':
img = simplify(img, hed_model_path)
elif a == 'face':
img = extract_face(img, target_encodings)
elif a == 'sss':
img = sss_processing.run_sss(img)
elif a == 'upsample':
img = esrgan_processing.upsample(img)
img = img.resize(
(int(img.width / 2), int(img.height / 2)),
resample=Image.BICUBIC) # go from 4x to 2x
elif a == 'none' or a == '':
pass
imgs1.append(img)
# save the images
for i, (img0, img1) in enumerate(zip(imgs0, imgs1)):
out_name = 'f%05d%s_%s.%s' % (idx_frame, '_%02d' % i if num_per > 1
else '', frame_name, save_ext)
is_train = training[num_per * k + i]
if save_mode == 'combined':
output_dir = trainA_dir if is_train else testA_dir
img2 = Image.new('RGB', (out_w * 2, out_h))
img2.paste(img1.convert('RGB'), (0, 0))
img2.paste(img0.convert('RGB'), (out_w, 0))
img2.save(os.path.join(output_dir, out_name), quality=97)
else:
if output_dir == 'None':
img1.convert('RGB').save(full_image_path, quality=97)
else:
outputA_dir = trainA_dir if is_train else testA_dir
img1.convert('RGB').save(os.path.join(
outputA_dir, out_name),
quality=97)
if save_mode == 'split':
outputB_dir = trainB_dir if is_train else testB_dir
img0.convert('RGB').save(os.path.join(
outputB_dir, out_name),
quality=97)
def main(args):
input_src, shuffle, max_num_images, min_w, min_h, max_w, max_h = args.input_src, args.shuffle, args.max_num_images, args.min_dim, args.min_dim, args.max_dim, args.max_dim
output_dir, out_w, out_h, pct_test, save_mode, save_ext = args.output_dir, args.w, args.h, args.pct_test, args.save_mode, args.save_ext
num_per, frac, frac_vary, max_ang_rot, max_stretch, centered = args.num_per, args.frac, args.frac_vary, args.max_ang_rot, args.max_stretch, args.centered
action, target_face_image, face_crop, face_crop_lerp, landmarks_path, hed_model_path = args.action, args.target_face_image, args.face_crop, args.face_crop_lerp, args.landmarks_path, args.hed_model_path
#os.system('rm -rf %s'%output_dir)
# get list of actions
actions = action.split(',')
if False in [a in allowable_actions for a in actions]:
raise Exception('one of your actions does not exist')
# initialize face_processing if needed
if 'face' in actions:
initialize_face_processing(landmarks_path)
target_encodings = get_encodings(target_face_image) if target_face_image else None
# initialize photosketch if needed
if 'sketch' in actions:
photosketch_processing.setup(args.photosketch_model_path)
# initialize esrgan if needed
if 'upsample' in actions:
esrgan_processing.setup(args.esrgan_model_path)
# initialize SSS if needed
if 'sss' in actions:
sss_processing.setup(args.sss_model_path)
# setup output directories
if output_dir != 'None':
trainA_dir, trainB_dir, testA_dir, testB_dir = setup_output_dirs(output_dir, save_mode, pct_test>0)
# initialize input
ext = os.path.splitext(input_src)[1]
is_movie = ext.lower() in ['.mp4','.mov','.avi']
if is_movie:
cap = cv2.VideoCapture(input_src)
fps = video.FPS().start()
num_images = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
pct_frames = list(np.linspace(0, 1, num_images))
all_frames = get_frame_indexes(max_num_images, num_images, shuffle)
else:
images = sorted([f for f in os.listdir(input_src) if os.path.isfile(os.path.join(input_src, f)) ])
num_images = len(images)
all_frames = get_frame_indexes(max_num_images, num_images, shuffle)
# training/test split
training = [1] * len(all_frames) * num_per
if pct_test > 0:
n_test = int(len(all_frames) * num_per * pct_test)
test_per = 1.0 / pct_test
test_idx = [int(test_per * (i+1) - 1) for i in range(n_test)]
for t in test_idx:
training[t] = 0
# iterate through each input
print("Iterating through %d input images" % len(all_frames))
for k, idx_frame in tqdm(enumerate(all_frames)):
if is_movie:
pct_frame = pct_frames[idx_frame]
frame = int(pct_frame * num_images)
cap.set(1, frame);
ret, img = cap.read()
frame_name = 'frame%06d' % frame
img = cv2pil(img)
else:
img_path = images[idx_frame]
frame_name = os.path.splitext(img_path)[0]
full_image_path = os.path.join(input_src, img_path)
img = Image.open(full_image_path).convert("RGB")
# skip images which are too small or too big
if img.width < min_w or img.height < min_h:
continue
if img.width > max_w or img.height > max_h:
continue
# first crop around face if requested
if face_crop is not None:
jx, jy, jw, jh = get_crop_around_face(img, target_encodings, out_w/out_h, face_crop, face_crop_lerp)
img = img.crop((jx, jy, jx + jw, jy + jh))
# preprocess/augment and produce input images
imgs0, imgs1 = augmentation(img, num_per, out_w, out_h, frac, frac_vary, max_ang_rot, max_stretch, centered), []
# process each input image to make output
for img0 in imgs0:
img = img0
for a in actions:
if a == 'segment':
img = segment(img)
elif a == 'colorize':
colors = [[255,255,255], [0,0,0], [127,0,0], [0, 0, 127], [0, 127, 0]]
img = quantize_colors(img, colors)
elif a == 'trace':
img = trace(img)
elif a == 'hed':
img = hed_processing.run_hed(img, hed_model_path)
elif a == 'sketch':
img = photosketch_processing.sketch(img)
elif a == 'simplify':
img = simplify(img, hed_model_path)
elif a == 'face':
img = extract_face(img, target_encodings)
elif a == 'sss':
img = sss_processing.run_sss(img)
elif a == 'upsample':
img = esrgan_processing.upsample(img)
img = img.resize((int(img.width/2), int(img.height/2)), resample=Image.BICUBIC) # go from 4x to 2x
elif a == 'none' or a == '':
pass
imgs1.append(img)
# save the images
for i, (img0, img1) in enumerate(zip(imgs0, imgs1)):
out_name = 'f%05d%s_%s.%s' % (idx_frame, '_%02d'%i if num_per>1 else '', frame_name, save_ext)
is_train = training[num_per * k + i]
if save_mode == 'combined':
output_dir = trainA_dir if is_train else testA_dir
img2 = Image.new('RGB', (out_w * 2, out_h))
img2.paste(img1.convert('RGB'), (0, 0))
img2.paste(img0.convert('RGB'), (out_w, 0))
img2.save(os.path.join(output_dir, out_name), quality=97)
else:
if output_dir == 'None':
img1.convert('RGB').save(full_image_path, quality=97)
else:
outputA_dir = trainA_dir if is_train else testA_dir
img1.convert('RGB').save(os.path.join(outputA_dir, out_name), quality=97)
if save_mode == 'split':
outputB_dir = trainB_dir if is_train else testB_dir
img0.convert('RGB').save(os.path.join(outputB_dir, out_name), quality=97)