Options Fix

app.py

@@ -30,7 +30,8 @@ import __init_paths
 from face_enhancement import FaceEnhancement 
 from face_colorization import FaceColorization 
 from face_inpainting import FaceInpainting 
-  
+from gradio_imageslider import ImageSlider
+
 def brush_stroke_mask(img, color=(255,255,255)):
     min_num_vertex = 8
     max_num_vertex = 28
@@ -87,49 +88,49 @@ def resize(image, width = 1024):
     image = cv2.resize(image, (int(height),int(width)))
     return image
 
-def inference(file, mode):
+def inference(file, mode, res_percentage, zoom, x_shift, y_shift):
 
-    im = cv2.imread(file, cv2.IMREAD_COLOR)
-    im = cv2.resize(im, (0,0), fx=2, fy=2)
-    faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=False)
-    img, orig_faces, enhanced_faces = faceenhancer.process(im)
-    cv2.imwrite(os.path.join("e.png"), img)
-    
+    im = cv2.resize(file, None, fx = (res_percentage/100), fy = (res_percentage/100))
     
     if mode == "enhance":
-        return os.path.join("e.png")
+        faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=False)
+        img, orig_faces, enhanced_faces = faceenhancer.process(im)
+        
     elif mode == "colorize":
         model = {'name':'GPEN-1024-Color', 'size':1024}
-        grayf = cv2.imread("e.png", cv2.IMREAD_GRAYSCALE)
+        if len(im.shape) == 3:
+            if im.shape[2] == 1:
+                grayf = im[:, :, 0]
+            elif im.shape[2] == 3:
+                grayf = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
+            elif im.shape[2] == 4:
+                grayf = cv2.cvtColor(im, cv2.COLOR_BGRA2GRAY)
+            grayf = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
+        
         grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR) # channel: 1->3
         facecolorizer = FaceColorization(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         colorf = facecolorizer.process(grayf)
-
-        colorf = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
-        cv2.imwrite(os.path.join("output.png"), colorf)
-        return os.path.join("output.png")
+        img = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
+        
     elif mode == "inpainting":
-        im1 = cv2.imread(file, cv2.IMREAD_COLOR)
         im2 = resize(im1, width = 1024)
         model = {'name':'GPEN-Inpainting-1024', 'size':1024}
         faceinpainter = FaceInpainting(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         im3 = np.asarray(brush_stroke_mask(Image.fromarray(im2)))
-        inpaint = faceinpainter.process(im3)
-
-        cv2.imwrite(os.path.join("output.png"), inpaint)
-        return os.path.join("output.png")
+        img = faceinpainter.process(im3)
+        
     elif mode == "selfie":
         model = {'name':'GPEN-BFR-2048', 'size':2048}
-        im = cv2.resize(im, (0,0), fx=2, fy=2)
         faceenhancer = FaceEnhancement(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         img, orig_faces, enhanced_faces = faceenhancer.process(im)
-        cv2.imwrite(os.path.join("output.png"), img)
-        return os.path.join("output.png")
+        
     else:
         faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=True)
         img, orig_faces, enhanced_faces = faceenhancer.process(im)
-        cv2.imwrite(os.path.join("output.png"), img)
-        return os.path.join("output.png")
+    
+    (in_img, out_img) = zoom_image(zoom, x_shift, y_shift, im, img)
+
+    return img, (in_img, out_img)
         
         
 title = "GPEN"
@@ -138,18 +139,60 @@ description = "Gradio demo for GAN Prior Embedded Network for Blind Face Restora
 article = "<p style='text-align: center;'><a href='https://arxiv.org/abs/2105.06070' target='_blank'>GAN Prior Embedded Network for Blind Face Restoration in the Wild</a> | <a href='https://github.com/yangxy/GPEN' target='_blank'>Github Repo</a></p><p style='text-align: center;'><img src='https://img.shields.io/badge/Hugging%20Face-Original%20demo-blue' alt='https://huggingface.co/spaces/akhaliq/GPEN' width='172' height='20' /></p>"
 
 
-gr.Interface(
-    inference, 
-    [gr.inputs.Image(type="filepath", label="Input"),gr.inputs.Radio(["enhance", "colorize", "inpainting", "selfie", "enhanced+background"], type="value", default="enhance", label="Type")], 
-    gr.outputs.Image(type="filepath", label="Output"),
-    title=title,
-    description=description,
-    article=article,
-    examples=[
-    ['enhance.png', 'enhance'],
-    ['color.png', 'colorize'],
-    ['inpainting.png', 'inpainting'],
-    ['selfie.png', 'selfie']
-    ],
-    enable_queue=True
-    ).launch()
+def zoom_image(zoom, x_shift, y_shift, input_img, output_img = None):
+    if output_img is None:
+        return None
+    
+    img = Image.fromarray(input_img)
+    out_img = Image.fromarray(output_img)
+    
+    img_w, img_h = img.size
+    zoom_factor = (100 - zoom) / 100
+    x_shift /= 100
+    y_shift /= 100
+    
+    zoom_w, zoom_h = int(img_w * zoom_factor), int(img_h * zoom_factor)
+    x_offset = int((img_w - zoom_w) * x_shift)
+    y_offset = int((img_h - zoom_h) * y_shift)
+    
+    crop_box = (x_offset, y_offset, x_offset + zoom_w, y_offset + zoom_h)
+    img = img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
+    out_img = out_img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
+
+    return (img, out_img)
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        input_img = gr.Image(label="Input Image")
+        output_img = gr.Image(label="Result")
+    
+    max_res = gr.Slider(1, 200, step=0.5, value=100, label="Output image Resolution Percentage (Higher% = longer processing time)")
+    restore_type = gr.Radio(["enhance", "colorize", "inpainting", "selfie", "enhanced+background"], value="enhance", type="value", label="Type")
+    
+    zoom = gr.Slider(0, 100, step=1, value=50, label="Zoom Percentage (0 = original size)")
+    x_shift = gr.Slider(0, 100, step=1, label="Horizontal shift Percentage (Before/After)")
+    y_shift = gr.Slider(0, 100, step=1, label="Vertical shift Percentage (Before/After)")
+    
+    run = gr.Button("Run")
+
+    with gr.Row():
+        before_after = ImageSlider(label="Before/After", type="pil", value=None)
+
+    run.click(
+        inference, 
+        inputs=[input_img, restore_type, max_res, zoom, x_shift, y_shift], 
+        outputs=[output_img, before_after]
+    )
+    
+    gr.Examples([
+        ['enhance.png', 'enhance', 100, 0, 0, 0],
+        ['color.png', 'colorize', 100, 0, 0, 0],
+        ['inpainting.png', 'inpainting', 100, 0, 0, 0],
+        ['selfie.png', 'selfie', 100, 0, 0, 0]
+    ], inputs=[input_img, restore_type, max_res, zoom, x_shift, y_shift])
+    
+    zoom.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img, output_img], outputs=[before_after])
+    x_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img, output_img], outputs=[before_after])
+    y_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img, output_img], outputs=[before_after])
+
+demo.launch()