Update app.py

app.py

@@ -71,10 +71,46 @@ def read_content(file_path):
         content = f.read()
     return content
     
+def get_image_enhancer(scale = 2, device='cuda:0'):
+    from basicsr.archs.rrdbnet_arch import RRDBNet
+    from realesrgan import RealESRGANer
+    from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+    from gfpgan import GFPGANer
+
+    realesrgan_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
+                    num_block=23, num_grow_ch=32, scale=4
+                  )
+    netscale = scale
+
+    model_realesrgan = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
+    upsampler = RealESRGANer(
+          scale=netscale,
+          model_path=model_realesrgan,
+          model=realesrgan_model,
+          tile=0,
+          tile_pad=10,
+          pre_pad=0,
+          half=False if device=='cpu' else True,
+          device=device
+    )
+
+    model_GFPGAN = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth'
+    img_enhancer = GFPGANer(
+          model_path=model_GFPGAN,
+          upscale=scale,
+          arch='clean',
+          channel_multiplier=2,
+          bg_upsampler=upsampler,
+          device=device
+      )
+    return img_enhancer
+
+image_enhancer = get_image_enhancer(scale = 1, device=device)
+        
 model = None
 
-def model_process(image, mask):
-    global model
+def model_process(image, mask, img_enhancer):
+    global model,image_enhancer
     
     if mask.shape[0] == image.shape[1] and mask.shape[1] == image.shape[0] and mask.shape[0] != mask.shape[1]:
         # rotate image
@@ -116,13 +152,13 @@ def model_process(image, mask):
     if config.sd_seed == -1:
         config.sd_seed = random.randint(1, 999999999)
 
-    print(f"Origin image shape_0_: {original_shape} / {size_limit}")
+    logger.info(f"Origin image shape_0_: {original_shape} / {size_limit}")
     image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
-    print(f"Resized image shape_1_: {image.shape}")
+    logger.info(f"Resized image shape_1_: {image.shape}")
     
-    print(f"mask image shape_0_: {mask.shape} / {type(mask)}")
+    logger.info(f"mask image shape_0_: {mask.shape} / {type(mask)}")
     mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
-    print(f"mask image shape_1_: {mask.shape} / {type(mask)}")
+    logger.info(f"mask image shape_1_: {mask.shape} / {type(mask)}")
 
     if model is None:
         return None
@@ -131,6 +167,19 @@ def model_process(image, mask):
     torch.cuda.empty_cache()
   
     image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, 'png')))
+    
+    if image_enhancer is not None and img_enhancer:
+        start = time.time()
+        input_img_rgb = np.array(image)
+        input_img_bgr = input_img_rgb[...,[2,1,0]]                        
+        _, _, enhance_img = image_enhancer.enhance(input_img_bgr, has_aligned=False,
+                                                only_center_face=False, paste_back=True)  
+        input_img_rgb = enhance_img[...,[2,1,0]]
+        img_enhance = Image.fromarray(np.uint8(input_img_rgb))                                                         
+        image = img_enhance  
+        log_info = f"image_enhancer_: {(time.time() - start) * 1000}ms, {res_np_img.shape} "
+        logger.info(log_info)
+         
     return  image # image
     
 model = ModelManager(
@@ -139,7 +188,9 @@ model = ModelManager(
     )
 
 image_type = 'pil' # filepath' 
-def predict(input):
+def predict(input, img_enhancer):
+    if input is None:
+        return None
     if image_type == 'filepath':
         # input: {'image': '/tmp/tmp8mn9xw93.png', 'mask': '/tmp/tmpn5ars4te.png'}
         origin_image_bytes = read_content(input["image"])
@@ -152,25 +203,25 @@ def predict(input):
         mask_pil = input['mask']
         image = np.array(image_pil)
         mask = np.array(mask_pil.convert("L"))
-    output = model_process(image, mask)
+    output = model_process(image, mask, img_enhancer)
     return output
 
 css = '''
 .container {max-width: 100%;margin: auto;padding-top: 1.5rem}
 .output-image, .input-image, .image-preview {height: 600px !important;object-fit: contain}
+#work-container {min-width: min(160px, 100%) !important;flex-grow: 0 !important}
 #image_upload{min-height:610px}
 #image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 620px}
 #image_output{margin: 0 auto; text-align: center;width:640px}
-#prompt-container{margin: 0 auto; text-align: center;width:200px;border-width:5px;border-color:#2c9748}
-#mask_radio .gr-form{background:transparent; border: none}
-#mask_radio .gr-form{background:transparent; border: none; color:#00ff00}
-#word_mask{margin-top: .75em !important}
-#word_mask textarea:disabled{opacity: 0.3}
+#erase-container{margin: 0 auto; text-align: center;width:150px;border-width:5px;border-color:#2c9748}
+#enhancer-checkbox{width:520px}
+#enhancer-tip{width:450px}
+#enhancer-tip-div{text-align: left}
+#prompt-container{margin: 0 auto; text-align: center;width:fit-content;min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;}
 .footer {margin-bottom: 45px;margin-top: 35px;text-align: center;border-bottom: 1px solid #e5e5e5}
 .footer>p {font-size: .8rem; display: inline-block; padding: 0 10px;transform: translateY(10px);background: white}
 .dark .footer {border-color: #303030}
 .dark .footer>p {background: #0b0f19}
-.acknowledgments h4{margin: 1.25em 0 .25em 0;font-weight: bold;font-size: 115%}
 #image_upload .touch-none{display: flex}
 @keyframes spin {
     from {
@@ -180,40 +231,30 @@ css = '''
         transform: rotate(360deg);
     }
 }
-#share-btn-container {
-    display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
-}
-#share-btn {
-    all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;
-}
-#share-btn * {
-    all: unset;
-}
-#share-btn-container div:nth-child(-n+2){
-    width: auto !important;
-    min-height: 0px !important;
-}
-#share-btn-container .wrap {
-    display: none !important;
-}
 '''
 
 image_blocks = gr.Blocks(css=css)
 with image_blocks as demo:
     with gr.Group():
-        with gr.Box():
-            with gr.Row():
+        with gr.Box(elem_id="work-container"):
+            with gr.Row(elem_id="input-container"):
                 with gr.Column():
-                    image = gr.Image(source='upload', elem_id="image_upload",tool='sketch', type=f'{image_type}', label="Upload").style(mobile_collapse=False)
-                    with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
-                        btn_in = gr.Button("Erase(↓)").style(
-                            margin=True,
-                            rounded=(True, True, True, True),
-                            full_width=True,
-                        )                
-            with gr.Row():             
+                    image = gr.Image(source='upload', elem_id="image_upload",tool='sketch', type=f'{image_type}', 
+                                     label="Upload(载入图片)", show_label=True).style(mobile_collapse=False)
+            with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
+                with gr.Column(elem_id="erase-container"):
+                    btn_erase = gr.Button(value = "Erase(擦除↓)",elem_id="erase_btn").style(
+                        margin=True,
+                        rounded=(True, True, True, True),
+                        full_width=True,
+                    ).style(width=100)   
+                with gr.Column(elem_id="enhancer-checkbox"):
+                    enhancer_label = 'Enhanced image(processing is very slow, please check only for blurred images)【增强图像(处理很慢，请仅针对模糊图像做勾选)】'
+                    img_enhancer = gr.Checkbox(label=enhancer_label).style(width=150) 
+            with gr.Row(elem_id="output-container"):             
                 with gr.Column():
-                    image_out = gr.Image(label="Output", elem_id="image_output", visible=True).style(width=640)
-            btn_in.click(fn=predict, inputs=[image], outputs=[image_out])
+                    image_out = gr.Image(label="Result", elem_id="image_output", visible=True).style(width=640)
+                    
+            btn_erase.click(fn=predict, inputs=[image, img_enhancer], outputs=[image_out])
             
-image_blocks.launch() 
+image_blocks.launch()