Create Streamlit demo

README.md

@@ -1,45 +1,20 @@
 ---
-title: Aot Gan Inpainting
+title: AOT-GAN Inpainting
 emoji: 🦀
 colorFrom: red
 colorTo: pink
 sdk: streamlit
 app_file: app.py
-pinned: false
+pinned: true
+sdk_version: 1.0.0
 ---
 
-# Configuration
+# AOT-GAN Inpainting
+This space demonstrates the [AOT-GAN for High-Resolution Image Inpainting](https://github.com/researchmm/AOT-GAN-for-Inpainting) developed
+by Yanhong Zeng, Jianlong Fu, Hongyang Chao, and Baining Guo. The GAN allows you to fill in large missing regions in high-resolution images.
 
-`title`: _string_
-Display title for the Space
-
-`emoji`: _string_
-Space emoji (emoji-only character allowed)
-
-`colorFrom`: _string_
-Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
-
-`colorTo`: _string_
-Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
-
-`sdk`: _string_
-Can be either `gradio`, `streamlit`, or `static`
-
-`sdk_version` : _string_
-Only applicable for `streamlit` SDK.  
-See [doc](https://huggingface.co/docs/hub/spaces) for more info on supported versions.
-
-`app_file`: _string_
-Path to your main application file (which contains either `gradio` or `streamlit` Python code, or `static` html code).
-Path is relative to the root of the repository.
-
-`models`: _List[string]_
-HF model IDs (like "gpt2" or "deepset/roberta-base-squad2") used in the Space.
-Will be parsed automatically from your code if not specified here.
-
-`datasets`: _List[string]_
-HF dataset IDs (like "common_voice" or "oscar-corpus/OSCAR-2109") used in the Space.
-Will be parsed automatically from your code if not specified here.
-
-`pinned`: _boolean_
-Whether the Space stays on top of your list.
+## Image Credit
+- `man.jpg` is sourced from https://github.com/researchmm/AOT-GAN-for-Inpainting
+- Photo by Ike louie Natividad from Pexels https://www.pexels.com/photo/woman-smiling-2709388/
+- Photo by Christina Morillo from Pexels https://www.pexels.com/photo/woman-smiling-at-the-camera-1181686/
+- Photo by Italo Melo from Pexels https://www.pexels.com/photo/man-wearing-blue-crew-neck-t-shirt-2379005/

app.py

@@ -0,0 +1,74 @@
+from PIL import Image
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+from torchvision.transforms import ToTensor
+import torch
+import numpy as np
+import cv2
+import aotgan.model.aotgan as net
+
+@st.cache
+def load_model(model_name):
+    model = net.InpaintGenerator.from_pretrained(model_name)
+    return model
+
+def postprocess(image):
+    image = torch.clamp(image, -1., 1.)
+    image = (image + 1) / 2.0 * 255.0
+    image = image.permute(1, 2, 0)
+    image = image.cpu().numpy().astype(np.uint8)
+    return image
+
+def infer(img, mask):
+    with torch.no_grad():
+        img_cv = cv2.resize(np.array(img), (512, 512))  # Fixing everything to 512 x 512 for this demo.
+        img_tensor = (ToTensor()(img_cv) * 2.0 - 1.0).unsqueeze(0)
+        mask_tensor = (ToTensor()(mask.astype(np.uint8))).unsqueeze(0)
+        masked_tensor = (img_tensor * (1 - mask_tensor).float()) + mask_tensor
+        pred_tensor = model(masked_tensor, mask_tensor)
+        comp_tensor = (pred_tensor * mask_tensor + img_tensor * (1 - mask_tensor))
+        comp_np = postprocess(comp_tensor[0])
+
+        return comp_np
+
+stroke_width = 8
+stroke_color = "#FFF"
+bg_color = "#000"
+bg_image = st.sidebar.file_uploader("Image:", type=["png", "jpg", "jpeg"])
+sample_bg_image = st.sidebar.radio('Sample Images', [
+    "man.png",
+    "pexels-ike-louie-natividad-2709388.jpg",
+    "pexels-christina-morillo-1181686.jpg",
+    "pexels-italo-melo-2379005.jpg",
+    "rainbow.jpeg",
+    "kitty.jpg",
+    "kitty_on_chair.jpeg",
+])
+drawing_mode = st.sidebar.selectbox(
+    "Drawing tool:", ("freedraw", "rect", "circle")
+)
+
+model_name = st.sidebar.selectbox(
+    "Select model:", ("NimaBoscarino/aot-gan-celebahq", "NimaBoscarino/aot-gan-places2")
+)
+model = load_model(model_name)
+
+bg_image = Image.open(bg_image) if bg_image else Image.open(f"./pictures/{sample_bg_image}")
+
+st.subheader("Draw on the image to erase features. The inpainted result will be generated and displayed below.")
+canvas_result = st_canvas(
+    fill_color="rgb(255, 255, 255)",
+    stroke_width=stroke_width,
+    stroke_color=stroke_color,
+    background_color=bg_color,
+    background_image=bg_image,
+    update_streamlit=True,
+    height=512,
+    width=512,
+    drawing_mode=drawing_mode,
+    key="canvas",
+)
+    
+if canvas_result.image_data is not None and bg_image and len(canvas_result.json_data["objects"]) > 0:
+    result = infer(bg_image, canvas_result.image_data[:, :, 3])
+    st.image(result)

requirements.txt

@@ -0,0 +1,7 @@
+git+https://github.com/NimaBoscarino/AOT-GAN-for-Inpainting.git
+streamlit_drawable_canvas
+opencv-python==4.5.1.48
+torch==1.8.1
+torchvision==0.9.1
+pillow==8.1.2
+transformers==4.15.0