Switch to Streamlit

README.md

@@ -1,10 +1,10 @@
 ---
-title: FcF Inpainting
+title: FcF-Inpainting
 emoji: 🪄 ✨ ✨
 colorFrom: blue
 colorTo: purple
-sdk: gradio
-sdk_version: 2.9.4
+sdk: streamlit
+sdk_version: 1.11.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -9,13 +9,9 @@ from PIL import Image
 import numpy as np
 import torch
 import legacy
-import paddlehub as hub
 import cv2
-
-u2net = hub.Module(name='U2Net')
-
-# gradio app imports
-import gradio as gr
+from streamlit_drawable_canvas import st_canvas
+import streamlit as st
 from torchvision.transforms import ToTensor, ToPILImage
 image_to_tensor = ToTensor()
 tensor_to_image = ToPILImage()
@@ -24,6 +20,17 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class_idx = None
 truncation_psi = 0.1
 
+title = "FcF-Inpainting"
+
+description = "<p style='color:royalblue; font-size: 14px; font-weight: w300;'>  \
+                [Note: The image and mask are resized to 512x512 before inpainting. The <span style='color:#E0B941;'>Run FcF-Inpainting</span> button will automatically appear after you draw a mask.] To use FcF-Inpainting: <br> \
+                (1) <span style='color:#E0B941;'>Upload an Image</span> or <span style='color:#E0B941;'> select a sample image on the left</span>. <br>  \
+                (2) Adjust the brush stroke width and <span style='color:#E0B941;'>draw the mask on the image</span>. You may also change the drawing tool on the sidebar. <br>\
+                (3) After drawing a mask, click the <span style='color:#E0B941;'>Run FcF-Inpainting</span> and witness the MAGIC! 🪄 ✨ ✨<br> \
+                (4) You may <span style='color:#E0B941;'>download/undo/redo/delete</span> the changes on the image using the options below the image box.</p>"
+
+article = "<p style='color: #E0B941; font-size: 16px; font-weight: w500; text-align: center'> <a style='color: #E0B941;' href='https://praeclarumjj3.github.io/fcf-inpainting/' target='_blank'>Project Page</a> | <a style='color: #E0B941;' href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'> Keys to Better Image Inpainting: Structure and Texture Go Hand in Hand</a> | <a style='color: #E0B941;' href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'>Github</a></p>"
+
 def create_model(network_pkl):
     print('Loading networks from "%s"...' % network_pkl)
     with dnnlib.util.open_url(network_pkl) as f:
@@ -48,11 +55,6 @@ def fcf_inpaint(G, org_img, erased_img, mask):
     comp_img = mask.to(device) * pred_img + (1 - mask).to(device) * org_img.to(device)
     return comp_img
 
-def show_images(img, width, height):
-    """ Display a batch of images inline. """
-    img = Image.fromarray(img)
-    img = img.resize((width, height))
-    return img
 
 def denorm(img):
     img = np.asarray(img[0].cpu(), dtype=np.float32).transpose(1, 2, 0)
@@ -64,38 +66,32 @@ def pil_to_numpy(pil_img: Image) -> Tuple[torch.Tensor, torch.Tensor]:
     img = np.array(pil_img)
     return torch.from_numpy(img)[None].permute(0, 3, 1, 2).float() / 127.5 - 1
 
-def inpaint(input_img, mask, option):
-    width, height = input_img.size
-
-    if option == "Automatic":
-        result = u2net.Segmentation(
-                images=[cv2.cvtColor(np.array(input_img), cv2.COLOR_RGB2BGR)],
-                paths=None,
-                batch_size=1,
-                input_size=320,
-                output_dir='output',
-                visualization=True)
-        mask = Image.fromarray(result[0]['mask'])
-        mask = mask.convert('L')
-    else:
-        mask = mask.resize((width,height))
+def process_mask(input_img, mask):
+    rgb = cv2.cvtColor(input_img, cv2.COLOR_RGBA2RGB)
+    mask = 255 - mask[:,:,3]
+    mask = (mask > 0) * 1
 
-    if width != 512 or height != 512:
-        input_img = input_img.resize((512, 512))
-        mask = mask.resize((512, 512))
-    
-    rgb = input_img.convert('RGB')
     rgb = np.array(rgb)
+    mask_tensor = torch.from_numpy(mask).to(torch.float32)
+    mask_tensor = mask_tensor.unsqueeze(0)
+    mask_tensor = mask_tensor.unsqueeze(0).to(device)
 
-    mask = np.array(mask)
-    if option == 'Manual':
-        mask = (mask[:, :, 0] == rgb[:, :, 0]) * (mask[:, :, 1] == rgb[:, :, 1]) * (mask[:, :, 2] == rgb[:, :, 2])
-        mask = 1. - mask.astype(np.float32) * 1.
-        kernel = np.ones((3, 3), np.uint8)
-        mask = cv2.dilate(mask, kernel) 
-        mask = mask * 255.
+    rgb = rgb.transpose(2,0,1)
+    rgb = torch.from_numpy(rgb.astype(np.float32)).unsqueeze(0)
+    rgb = (rgb.to(torch.float32) / 127.5 - 1).to(device)
+    rgb_erased = rgb.clone()
+    rgb_erased = rgb_erased * (1 - mask_tensor) # erase rgb
+    rgb_erased = rgb_erased.to(torch.float32)
+    
+    rgb_erased = denorm(rgb_erased)
+    return rgb_erased
 
-    mask /= 255.
+def inpaint(input_img, mask, model):
+    rgb = cv2.cvtColor(input_img, cv2.COLOR_RGBA2RGB)
+    mask = 255 - mask[:,:,3]
+    mask = (mask > 0) * 1
+
+    rgb = np.array(rgb)
     mask_tensor = torch.from_numpy(mask).to(torch.float32)
     mask_tensor = mask_tensor.unsqueeze(0)
     mask_tensor = mask_tensor.unsqueeze(0).to(device)
@@ -107,60 +103,93 @@ def inpaint(input_img, mask, option):
     rgb_erased = rgb_erased * (1 - mask_tensor) # erase rgb
     rgb_erased = rgb_erased.to(torch.float32)
     
-    model = create_model("models/places_512.pkl")
     comp_img = fcf_inpaint(G=model, org_img=rgb.to(torch.float32), erased_img=rgb_erased.to(torch.float32), mask=mask_tensor.to(torch.float32))
     rgb_erased = denorm(rgb_erased)
     comp_img = denorm(comp_img)
+    return comp_img
 
-    return show_images(rgb_erased, width, height), show_images(comp_img, width, height)
-
-gradio_inputs = [gr.inputs.Image(type='pil',
-                                 tool="editor",
-                                 label="Image"),
-                # gr.inputs.Image(type='pil',source="canvas", label="Mask", invert_colors=True),
-                gr.inputs.Image(type='pil',
-                                 tool="editor",
-                                 label="Mask"),
-                gr.inputs.Radio(choices=["Automatic", "Manual"], type="value", default="Manual", label="Masking Choice")
-                ]
-
-gradio_outputs = [gr.outputs.Image(label='Image with Hole'),
-                 gr.outputs.Image(label='Inpainted Image')]
-
-examples = [['test_512/person512.png', 'test_512/person512.png', 'Automatic'],
-            ['test_512/a_org.png', 'test_512/a_overlay.png', 'Manual'],
-            ['test_512/f_org.png', 'test_512/f_overlay.png', 'Manual'],
-            ['test_512/g_org.png', 'test_512/g_overlay.png', 'Manual'],
-            ['test_512/h_org.png', 'test_512/h_overlay.png', 'Manual'],
-            ['test_512/i_org.png', 'test_512/i_overlay.png', 'Manual'],
-            ['test_512/b_org.png', 'test_512/b_overlay.png', 'Manual'],
-            ['test_512/c_org.png', 'test_512/c_overlay.png', 'Manual'],
-            ['test_512/d_org.png', 'test_512/d_overlay.png', 'Manual'],
-            ['test_512/e_org.png', 'test_512/e_overlay.png', 'Manual']]
-
-title = "FcF-Inpainting"
+def run_app(model):
+    
+    if "button_id" not in st.session_state:
+        st.session_state["button_id"] = ""
+    if "color_to_label" not in st.session_state:
+        st.session_state["color_to_label"] = {}
+    image_inpainting(model)
+
+    with st.sidebar:
+        st.markdown("---")
+
+def image_inpainting(model):
+    if 'reuse_image' not in st.session_state:
+        st.session_state.reuse_image = None
+    
+    st.title(title)
+    st.markdown(article, unsafe_allow_html=True)
+    st.markdown(description, unsafe_allow_html=True)
 
-description = "<p style='color:royalblue; font-weight: w300;'>  \
-                [Note: Queue time may take up to 20 seconds! The image and mask are resized to 512x512 before inpainting.] To use FcF-Inpainting: <br> \
-                (1) <span style='color:#E0B941;'>Upload the Same Image to both</span> input boxes (Image and Mask) below. <br>  \
-                (2a) <span style='color:#E0B941;'>Manual Option:</span> The TUI Image Editor used by gradio <a style='color: #E0B941;' href='https://github.com/gradio-app/gradio/issues/1810' target='_blank'>changes the image when saved</a>. We compute the mask after comparing the two inputs. Therefore, we need to save both inputs: <br> \
-                &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; - <span style='color:#E0B941;'>Image:</span> Click on the edit button on the top-right and save without making any changes. <br> \
-                &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; - <span style='color:#E0B941;'>Mask:</span> Draw a mask (hole) using the brush (click on the edit button in the top right of the Mask View and select the draw option). <br> \
-                (2b) <span style='color:#E0B941;'>Automatic Option:</span> This option will generate a mask using a pretrained U2Net model. <br>  \
-                (3) Click on <span style='color:#E0B941;'>Submit</span> and witness the MAGIC! 🪄 ✨ ✨</p>"
-
-article = "<p style='color: #E0B941; text-align: center'> <a style='color: #E0B941;' href='https://praeclarumjj3.github.io/fcf-inpainting/' target='_blank'>Project Page</a> | <a style='color: #E0B941;' href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'> Keys to Better Image Inpainting: Structure and Texture Go Hand in Hand</a> | <a style='color: #E0B941;' href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'>Github Repo</a></p>"
-
-css = ".image-preview {height: 32rem; width: auto;} .output-image {height: 32rem; width: auto;} .panel-buttons { display: flex; flex-direction: row;}"
-
-iface = gr.Interface(fn=inpaint, inputs=gradio_inputs,
-                     outputs=gradio_outputs,
-                     css=css,
-                     layout="vertical",
-                     theme="dark-huggingface",
-                     examples_per_page=5,
-                     thumbnail="fcf_gan.png",
-                     allow_flagging="never",
-                     examples=examples, title=title,
-                     description=description, article=article)
-iface.launch(enable_queue=True)
+    image = st.sidebar.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
+    
+    sample_image = st.sidebar.radio('Choose a Sample Image', [
+            'scene-background.png',
+            'fence-background.png',
+            'bench.png',
+            'house.png',
+            'landscape.png',
+            'truck.png',
+            'scenery.png',
+            'grass-texture.png',
+            'mapview-texture.png',
+        ])
+    
+    drawing_mode = st.sidebar.selectbox(
+    "Drawing tool:", ("freedraw", "line")
+)
+
+    image = Image.open(image).convert("RGBA") if image else Image.open(f"./test_512/{sample_image}").convert("RGBA")        
+    image = image.resize((512, 512))
+    width, height = image.size
+    stroke_width = st.sidebar.slider("Stroke width: ", 1, 100, 20)
+
+    canvas_result = st_canvas(
+        stroke_color="rgba(255, 0, 255, 0.8)",
+        stroke_width=stroke_width,
+        background_image=image,
+        height=height,
+        width=width,
+        drawing_mode=drawing_mode,
+        key="canvas",
+    )
+    if canvas_result.image_data is not None and image and len(canvas_result.json_data["objects"]) > 0:
+        
+        im = canvas_result.image_data.copy()
+        background = np.where(
+            (im[:, :, 0] == 0) & 
+            (im[:, :, 1] == 0) & 
+            (im[:, :, 2] == 0)
+        )
+        drawing = np.where(
+            (im[:, :, 0] == 255) & 
+            (im[:, :, 1] == 0) & 
+            (im[:, :, 2] == 255)
+        )
+        im[background]=[0,0,0,255]
+        im[drawing]=[0,0,0,0] #RGBA
+        if st.button('Run FcF-Inpainting'):
+            col1, col2 = st.columns([1,1])
+            with col1:
+                # if st.button('Show Image with Holes'):
+                st.write("Masked Image")
+                mask_show = process_mask(np.array(image), np.array(im))
+                st.image(mask_show)
+            with col2:
+                st.write("Inpainted Image")
+                inpainted_img = inpaint(np.array(image), np.array(im), model)
+                st.image(inpainted_img)
+
+if __name__ == "__main__":
+    st.set_page_config(
+        page_title="FcF-Inpainting", page_icon=":sparkles:"
+    )
+    st.sidebar.subheader("Configuration")
+    model = create_model("models/places_512.pkl")
+    run_app(model)

requirements.txt

@@ -1,7 +1,5 @@
 icecream==2.1.0
 opencv-python-headless
-paddlepaddle
-paddlehub
 psutil==5.8.0
 click 
 requests
@@ -22,6 +20,7 @@ pydrive2
 pandas
 easydict
 kornia==0.5.0
-gradio==2.9.4
 ipython
-Jinja2
+Jinja2
+streamlit-drawable-canvas
+streamlit==1.11.0