Skip to content

Commit

Permalink
minor updates
Browse files Browse the repository at this point in the history 
1. FaceSwap_GAN_v2.2_train_test.ipynb: Make code for single image tranformation more compact.

2. data_loader\data_augmentation.py: Optimize roi area for data augmentation.
  • Loading branch information
@shaoanlu
shaoanlu committed Jan 8, 2019
1 parent 8bb0559 commit 1d604c5
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 87 deletions.
131 changes: 45 additions & 86 deletions FaceSwap_GAN_v2.2_train_test.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,17 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 1,
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Using TensorFlow backend.\n"
]
}
],
"source": [
"from keras.layers import *\n",
"import keras.backend as K\n",
Expand All @@ -21,7 +29,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -47,7 +55,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -57,7 +65,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -92,7 +100,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -725,7 +733,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -735,7 +743,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -745,84 +753,18 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"def shift_n_scale(src_img, tar_img):\n",
" \"\"\"\n",
" A color correction method\n",
" \"\"\"\n",
" src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2XYZ)\n",
" tar_img = cv2.cvtColor(tar_img, cv2.COLOR_BGR2XYZ) \n",
" mt = np.mean(tar_img, axis=(0,1))\n",
" st = np.std(tar_img, axis=(0,1))\n",
" ms = np.mean(src_img, axis=(0,1))\n",
" ss = np.std(src_img, axis=(0,1)) \n",
" if ss.any() <= 1e-7: return src_img \n",
" result = st * (src_img.astype(np.float32) - ms) / (ss+1e-7) + mt\n",
" result = (255.0 / result.max() * result).astype(np.float32) \n",
" result = cv2.cvtColor(result.astype(np.uint8), cv2.COLOR_XYZ2BGR)\n",
" return result\n",
"\n",
"def transform_face(inp_img, direction=\"AtoB\", roi_coef=15, color_correction=True):\n",
" \"\"\"\n",
" Parameters:\n",
" inp_img: A RGB face image of any size.\n",
" direction: A string that is either AtoB or BtoA\n",
" roi_coef: A coefficient that affects the cropped center area\n",
" color_correction: boolean, whether use color correction or not\n",
" Returns:\n",
" result_img: A RGB swapped face image after masking.\n",
" result_mask: The alpha mask which corresponds to the result_img.\n",
" \"\"\"\n",
" def get_feather_edges_mask(img, roi_coef):\n",
" img_size = img.shape\n",
" mask = np.zeros_like(img)\n",
" x_, y_ = img_size[0]//roi_coef, img_size[1]//roi_coef\n",
" mask[x_:-x_, y_:-y_,:] = 255\n",
" mask = cv2.GaussianBlur(mask,(15,15),10).astype(np.float32) / 255\n",
" return mask \n",
"\n",
" if direction == \"AtoB\":\n",
" path_func = model.path_abgr_B\n",
" elif direction == \"BtoA\":\n",
" path_func = model.path_abgr_A\n",
" else:\n",
" raise ValueError(f\"direction should be either AtoB or BtoA, recieved {direction}.\")\n",
"\n",
" # pre-process input image\n",
" img_bgr = cv2.cvtColor(inp_img, cv2.COLOR_RGB2BGR)\n",
" input_size = img_bgr.shape \n",
" roi_x, roi_y = input_size[0]//roi_coef, input_size[1]//roi_coef\n",
" roi = img_bgr[roi_x:-roi_x, roi_y:-roi_y,:] # BGR, [0, 255] \n",
" roi_size = roi.shape\n",
" ae_input = cv2.resize(roi, (RESOLUTION,RESOLUTION)) / 255. * 2 - 1 # BGR, [-1, 1] \n",
"\n",
" # post-process transformed roi image\n",
" ae_output = np.squeeze(np.array([path_func([[ae_input]])]))\n",
" ae_output_a = ae_output[:,:,0] * 255\n",
" ae_output_a = cv2.resize(ae_output_a, (roi_size[1],roi_size[0]))[...,np.newaxis]\n",
" ae_output_bgr = np.clip( (ae_output[:,:,1:] + 1) * 255 / 2, 0, 255)\n",
" ae_output_bgr = cv2.resize(ae_output_bgr, (roi_size[1],roi_size[0]))\n",
" ae_output_masked = (ae_output_a/255 * ae_output_bgr + (1 - ae_output_a/255) * roi).astype('uint8') # BGR, [0, 255]\n",
" if color_correction:\n",
" ae_output_masked = shift_n_scale(ae_output_masked, roi)\n",
"\n",
" # merge transformed output back to input image\n",
" blend_mask = get_feather_edges_mask(roi, roi_coef) \n",
" blended_img = blend_mask * ae_output_masked + (1 - blend_mask) * roi\n",
" result = img_bgr\n",
" result[roi_x:-roi_x, roi_y:-roi_y,:] = blended_img \n",
" result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) \n",
" result_alpha = np.zeros_like(img_bgr)\n",
" result_alpha[roi_x:-roi_x, roi_y:-roi_y,:] = blend_mask * ae_output_a \n",
" return result, result_alpha"
"from converter.face_transformer import FaceTransformer\n",
"ftrans = FaceTransformer()\n",
"ftrans.set_model(model)"
]
},
{
"cell_type": "code",
"execution_count": 20,
"execution_count": 70,
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -873,21 +815,28 @@
},
{
"cell_type": "code",
"execution_count": 22,
"execution_count": 73,
"metadata": {},
"outputs": [],
"source": [
"# Transform detected face\n",
"result_img, result_mask = transform_face(aligned_det_face_im, direction=\"AtoB\", roi_coef=15)\n",
"result_img, result_rgb, result_mask = ftrans.transform(\n",
" aligned_det_face_im, \n",
" direction=\"AtoB\", \n",
" roi_coverage=0.93,\n",
" color_correction=\"adain_xyz\",\n",
" IMAGE_SHAPE=(RESOLUTION, RESOLUTION, 3)\n",
" )\n",
"try:\n",
" result_img = landmarks_match_mtcnn(result_img, tar_landmarks, src_landmarks)\n",
" result_rgb = landmarks_match_mtcnn(result_rgb, tar_landmarks, src_landmarks)\n",
" result_mask = landmarks_match_mtcnn(result_mask, tar_landmarks, src_landmarks)\n",
"except:\n",
" print(\"An error occured during face alignment.\")\n",
" pass\n",
"\n",
"result_input_img = input_img.copy()\n",
"result_input_img[int(x0):int(x1),int(y0):int(y1),:] = result_mask.astype(np.float32)/255*result_img +\\\n",
"result_input_img[int(x0):int(x1),int(y0):int(y1),:] = result_mask.astype(np.float32)/255*result_rgb +\\\n",
"(1-result_mask.astype(np.float32)/255)*result_input_img[int(x0):int(x1),int(y0):int(y1),:]"
]
},
Expand All @@ -897,8 +846,18 @@
"metadata": {},
"outputs": [],
"source": [
"# Display transformed image\n",
"plt.imshow(result_img)"
"# Show result face\n",
"plt.imshow(result_input_img[int(x0):int(x1),int(y0):int(y1),:])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Show transformed image before masking\n",
"plt.imshow(result_rgb)"
]
},
{
Expand All @@ -907,7 +866,7 @@
"metadata": {},
"outputs": [],
"source": [
"# Display alpha mask\n",
"# Show alpha mask\n",
"plt.imshow(result_mask[..., 0])"
]
},
Expand Down Expand Up @@ -952,7 +911,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.4"
"version": "3.6.6"
}
},
"nbformat": 4,
Expand Down
2 changes: 1 addition & 1 deletion data_loader/data_augmentation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def random_warp_rev(image, res=64):
interp_slice = slice(interp_param//10,9*interp_param//10)
dst_pnts_slice = slice(0,65*res_scale,16*res_scale)

rand_coverage = np.random.randint(20) + 85 # random warping coverage
rand_coverage = np.random.randint(20) + 78 # random warping coverage
rand_scale = np.random.uniform(5., 6.2) # random warping scale

range_ = np.linspace(128-rand_coverage, 128+rand_coverage, 5)
Expand Down

0 comments on commit 1d604c5

Please sign in to comment.