diff --git a/modelcenter/PP-Matting/APP/app.py b/modelcenter/PP-Matting/APP/app.py
index 26f4aabcfb9e078cc460c9ce9211f7135706c747..6318adb8046ef5e9984f0a827e6bbf14b6f588b4 100644
--- a/modelcenter/PP-Matting/APP/app.py
+++ b/modelcenter/PP-Matting/APP/app.py
@@ -1,15 +1,111 @@
import gradio as gr
import numpy as np
+import cv2
import utils
from predict import build_predictor
-IMAGE_DEMO = "./images/idphoto.jpg"
+ID_PHOTO_IMAGE_DEMO = "./images/idphoto.jpg"
+
+MC_PHOTO_BG = "Red"
+MC_PHOTO_SIZE = (413, 626)
+MC_PHOTO_WIFE_IMAGE_DEMO = "./images/wife.jpg"
+MC_PHOTO_HUSBAND_IMAGE_DEMO = "./images/husband.jpg"
+
predictor = build_predictor()
sizes_play = utils.size_play()
-def get_output(img, size, bg, download_size):
+def crop(img, alpha, thr=0.001):
+ """
+ Crop the image and alpha according to alpha mask.
+
+ Args:
+ img(numpy:ndarray): The image array.
+ alpha(numpy:ndarray): The alpha corresponding to the image.
+ thr(float): The threshold used to generate the alpha mask.
+
+ Returns:
+ img_cropped(numpy:ndarray): Image after crop.
+ alpha_cropped(numpy:ndarray): alpha after crop.
+ """
+ _, mask = cv2.threshold(alpha, 0, 255, cv2.THRESH_BINARY, thr)
+ cnt, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+ x, y, w, h = cv2.boundingRect(cnt[0])
+ return img[y: y + h, x: x + w], alpha[y: y + h, x: x + w]
+
+
+def copy_img_to_bg(img, alpha, target_h, target_x, bg_color, bg_size):
+ """
+ Copy img and alpha to the corresponding background respectively.
+
+ Args:
+ img(numpy:ndarray): The image array.
+ alpha(numpy:ndarray): The alpha corresponding to the image.
+ target_h(int): The target height of img in background.
+ target_x(int): The target position of the img's x-center in the background.
+ bg_color(str): The color of background, the options are "White"、"Blue"、"Red".
+ bg_size(tuple): The size of background.
+
+ Returns:
+ res_img(numpy:ndarray): Result after copy img to background.
+ res_alpha(numpy:ndarray): Result after copy alpha to background.
+ """
+ bg_h, bg_w = bg_size
+ img_h, img_w, _ = img.shape
+ r = 1.0 * img_h / target_h
+ target_w = int(img_w / r)
+
+ res_img = np.ones((bg_h, bg_w, 3)) * utils.COLOR_MAP[bg_color]
+ res_alpha = np.zeros((bg_h, bg_w))
+
+ img = cv2.resize(img, (target_w, target_h), interpolation=cv2.INTER_AREA)
+ alpha = cv2.resize(alpha, (target_w, target_h), interpolation=cv2.INTER_AREA)
+
+ x1 = int(max(0, target_x - target_w / 2))
+ y1 = int(max(0, bg_h - target_h))
+ x2 = int(min(bg_w, x1+target_w))
+ y2 = int(min(bg_h, y1+target_h))
+ act_w = x2 - x1
+ act_h = y2 - y1
+ res_img[y1: y2, x1: x2] = img[:act_h, :act_w]
+ res_alpha[y1: y2, x1: x2] = alpha[:act_h, :act_w]
+
+ return res_img, res_alpha
+
+
+def get_mc_photo_app_output(wife, husband):
+ """
+ Composite marriage certificate photo
+
+ Args:
+ wife(numpy:ndarray): The wife image array.
+ husband(numpy:ndarray): The husband image array.
+
+ Returns:
+ res(numpy:ndarray): The composited marriage certificate photo.
+ res_download(str): The path of res.
+ """
+ alpha_wife = predictor.run(wife)
+ alpha_husband = predictor.run(husband)
+
+ _, wife_fg = utils.bg_replace(wife, alpha_wife, bg_name=MC_PHOTO_BG)
+ husband, _ = utils.bg_replace(husband, alpha_husband, bg_name=MC_PHOTO_BG)
+
+ wife_fg, alpha_wife = crop(wife_fg, alpha_wife)
+ husband, alpha_husband = crop(husband, alpha_husband)
+
+ wife_fg, alpha_wife = copy_img_to_bg(wife_fg, alpha_wife, 338, 225, MC_PHOTO_BG, MC_PHOTO_SIZE)
+ husband, alpha_husband = copy_img_to_bg(husband, alpha_husband, 375, 401, MC_PHOTO_BG, MC_PHOTO_SIZE)
+
+ alpha_wife = alpha_wife[:, :, None] / 255.0
+ res = (wife_fg * alpha_wife + (1 - alpha_wife) * husband).astype(np.uint8)
+
+ res_download = utils.download(res, "Large")
+ return res, res_download
+
+
+def get_id_photo_output(img, size, bg, download_size):
"""
Get the special size and background photo.
@@ -21,7 +117,7 @@ def get_output(img, size, bg, download_size):
"""
alpha = predictor.run(img)
- res = utils.bg_replace(img, alpha, bg_name=bg)
+ res, _ = utils.bg_replace(img, alpha, bg_name=bg)
size_index = sizes_play.index(size)
res = utils.adjust_size(res, size_index)
@@ -29,49 +125,79 @@ def get_output(img, size, bg, download_size):
return res, res_download
-def download(img, size):
- utils.download(img, size)
- return None
+def clear_id_photo_all():
+ utils.delete_result()
+ return None, None, sizes_play[0], 'White', 'Large', None
-with gr.Blocks() as demo:
- gr.Markdown("""# ID Photo DIY""")
+def clear_mc_photo_all():
+ utils.delete_result()
+ return None, None, None, None
+
- img_in = gr.Image(value=IMAGE_DEMO, label="Input image")
+with gr.Blocks() as demo:
+ gr.Markdown("""# ID and MC(Marriage Certificate) Photo DIY""")
gr.Markdown(
- """Tips: Please upload photos with good posture, center portrait, crown free, no jewelry, ears and eyebrows exposed."""
+ """Tips: Please upload photos with good posture, center portrait,
+ crown free, no jewelry, ears and eyebrows exposed."""
)
- with gr.Row():
- size = gr.Dropdown(sizes_play, label="Sizes", value=sizes_play[0])
- bg = gr.Radio(
- ["White", "Red", "Blue"], label="Background color", value='White')
- download_size = gr.Radio(
- ["Small", "Middle", "Large"],
- label="File size (affects image quality)",
- value='Large',
- interactive=True)
-
- with gr.Row():
- btn1 = gr.Button("Clear")
- btn2 = gr.Button("Submit")
-
- img_out = gr.Image(
- label="Output image", interactive=False).style(height=300)
- f1 = gr.File(label='Image download').style(height=50)
- with gr.Row():
- gr.Markdown(
- """This application is supported by [PaddleSeg](https://github.com/PaddlePaddle/PaddleSeg).
- If you have any questions or feature requists, welcome to raise issues on [GitHub](https://github.com/PaddlePaddle/PaddleSeg/issues). BTW, a star is a great encouragement for us, thanks! ^_^"""
- )
-
- btn2.click(
- fn=get_output,
- inputs=[img_in, size, bg, download_size],
- outputs=[img_out, f1])
- btn1.click(
- fn=utils.clear_all,
+ with gr.Tab("IDPhoto"):
+ id_photo_img_in = gr.Image(value=ID_PHOTO_IMAGE_DEMO, label="Input image")
+ with gr.Row():
+ id_photo_size = gr.Dropdown(sizes_play, label="Sizes", value=sizes_play[0])
+ id_photo_bg = gr.Radio(
+ ["White", "Red", "Blue"], label="Background color", value='White')
+ id_photo_download_size = gr.Radio(
+ ["Small", "Middle", "Large"],
+ label="File size (affects image quality)",
+ value='Large',
+ interactive=True)
+
+ with gr.Row():
+ id_photo_clear_btn = gr.Button("Clear")
+ id_photo_submit_btn = gr.Button("Submit")
+
+ id_photo_img_out = gr.Image(label="Output image", interactive=False).style(height=300)
+ id_photo_downloaded_img = gr.File(label='Image download').style(height=50)
+
+ id_photo_clear_btn.click(
+ fn=clear_id_photo_all,
+ inputs=None,
+ outputs=[id_photo_img_in, id_photo_img_out, id_photo_size, id_photo_bg,
+ id_photo_download_size, id_photo_downloaded_img])
+ id_photo_submit_btn.click(
+ fn=get_id_photo_output,
+ inputs=[id_photo_img_in, id_photo_size, id_photo_bg, id_photo_download_size],
+ outputs=[id_photo_img_out, id_photo_downloaded_img])
+
+ with gr.Tab("MCPhoto"):
+ with gr.Row():
+ mc_photo_img_wife = gr.Image(value=MC_PHOTO_WIFE_IMAGE_DEMO, label="Wife", interactive=True)
+ mc_photo_img_husband = gr.Image(value=MC_PHOTO_HUSBAND_IMAGE_DEMO, label="Husband", interactive=True)
+
+ with gr.Row():
+ mc_photo_clear_button = gr.Button("Clear")
+ mc_photo_submit_button = gr.Button("Submit")
+
+ mc_photo_img_out = gr.Image(label="Output image", interactive=False).style(height=300)
+ mc_photo_download_img = gr.File(label='Image download').style(height=50)
+
+ mc_photo_clear_button.click(
+ fn=clear_mc_photo_all,
inputs=None,
- outputs=[img_in, img_out, size, bg, download_size, f1])
+ outputs=[mc_photo_img_wife, mc_photo_img_husband, mc_photo_img_out, mc_photo_download_img])
+ mc_photo_submit_button.click(
+ fn=get_mc_photo_app_output,
+ inputs=[mc_photo_img_wife, mc_photo_img_husband],
+ outputs=[mc_photo_img_out, mc_photo_download_img])
+
+ gr.Markdown(
+ """This application is supported by
+ [PaddleSeg](https://github.com/PaddlePaddle/PaddleSeg).
+ If you have any question or feature request,
+ welcome to raise issues on [GitHub](https://github.com/PaddlePaddle/PaddleSeg/issues).
+ BTW, a star is a great encouragement for us, thanks! ^_^"""
+ )
gr.Button.style(1)
diff --git a/modelcenter/PP-Matting/APP/images/husband.jpg b/modelcenter/PP-Matting/APP/images/husband.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..e79a52d3b9753ba0324fc774a935b6f156eb3d52
Binary files /dev/null and b/modelcenter/PP-Matting/APP/images/husband.jpg differ
diff --git a/modelcenter/PP-Matting/APP/images/wife.jpg b/modelcenter/PP-Matting/APP/images/wife.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..652a1a5eefb2c5e6419eb6bbfeb582002dbcb705
Binary files /dev/null and b/modelcenter/PP-Matting/APP/images/wife.jpg differ
diff --git a/modelcenter/PP-Matting/APP/utils.py b/modelcenter/PP-Matting/APP/utils.py
index b616d3e4de7fde6e6df31b1f62f53937ee3e2999..8e0147bb9c81b0fabedc3c27301a741b00ab828a 100644
--- a/modelcenter/PP-Matting/APP/utils.py
+++ b/modelcenter/PP-Matting/APP/utils.py
@@ -56,19 +56,15 @@ COLOR_MAP = {
# jpg compress ratio
SAVE_SIZE = {'Small': 50, 'Middle': 75, 'Large': 95}
+CACHE_DIR = ".temp"
+
def delete_result():
"""clear old result in `.temp`"""
- root = '.temp'
- results = sorted(os.listdir(root))
+ results = sorted(os.listdir(CACHE_DIR))
for res in results:
- if int(time.time()) - int(os.path.splitext(res)[0]) > 900:
- os.remove(os.path.join(root, res))
-
-
-def clear_all():
- delete_result()
- return None, None, size_play()[0], 'White', 'Large', None
+ if int(time.time()) - int(os.path.splitext(res)[0]) > 10000:
+ os.remove(os.path.join(CACHE_DIR, res))
def size_play():
@@ -86,10 +82,10 @@ def bg_replace(img, alpha, bg_name):
bg = COLOR_MAP[bg_name]
bg = np.array(bg)[None, None, :]
alpha = alpha / 255.
- pymatting.estimate_foreground_ml(img / 255., alpha) * 255
+ fg = pymatting.estimate_foreground_ml(img / 255., alpha, return_background=False) * 255
alpha = alpha[:, :, None]
- res = alpha * img + (1 - alpha) * bg
- return res.astype('uint8')
+ res = alpha * fg + (1 - alpha) * bg
+ return res.astype('uint8'), fg.astype('uint8')
def adjust_size(img, size_index):
@@ -121,16 +117,15 @@ def adjust_size(img, size_index):
def download(img, size):
q = SAVE_SIZE[size]
+ if not os.path.exists(CACHE_DIR):
+ os.makedirs(CACHE_DIR)
while True:
name = str(int(time.time()))
- tmp_name = './.temp/' + name + '.jpg'
+ tmp_name = os.path.join(CACHE_DIR, name + '.jpg')
if not os.path.exists(tmp_name):
break
else:
time.sleep(1)
- dir_name = os.path.dirname(tmp_name)
- if not os.path.exists(dir_name):
- os.makedirs(dir_name)
im = Image.fromarray(img)
im.save(tmp_name, 'jpeg', quality=q, dpi=(300, 300))