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))