add demo

Author: zengxianyu

README.md

@@ -61,11 +61,10 @@ you may modify the training script to use different settings, e.g., batch size,
 
 ### Web APP
 <img src="https://s3.ax1x.com/2020/11/27/DrVLs1.png" width=300>
-For your convinience of visualization and evaluation, I provide an inpainting APP where you can interact with the inpainting model in a browser, to open a photo and draw area to remove. To use the web app, these additional packages are required: 
 
-```flask```, ```requests```, ```pillow```
+To use the web app, these additional packages are required: 
 
-Then execute the following:
+```flask```, ```requests```, ```pillow```
 
 With GPU:
 ```

demo.py

@@ -0,0 +1,137 @@
+import pdb
+import cv2
+import os
+from collections import OrderedDict
+
+import numpy as np
+from werkzeug.utils import secure_filename
+from flask import Flask, url_for, render_template, request, redirect, send_from_directory
+from PIL import Image
+import base64
+import io
+import random
+
+
+from options.test_options import TestOptions
+import models
+import torch
+
+opt = TestOptions().parse()
+model = models.create_model(opt)
+model.eval()
+
+max_size = 256
+max_num_examples = 200
+UPLOAD_FOLDER = 'static/images'
+app = Flask(__name__)
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+ALLOWED_EXTENSIONS = set(['png', 'jpg', 'JPG', 'jpeg', 'bmp'])
+def allowed_file(filename):
+    return '.' in filename and \
+           filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS
+
+port = opt.port
+filelist = "./static/images/example.txt"
+with open(filelist, "r") as f:
+    list_examples = f.readlines()
+list_examples = [n.strip("\n") for n in list_examples]
+
+def process_image(img, mask, name, opt, save_to_input=True):
+    img =img.convert("RGB")
+    img_raw = np.array(img)
+    w_raw, h_raw = img.size
+    h_t, w_t = h_raw//8*8, w_raw//8*8
+
+    img = img.resize((w_t, h_t))
+    img = np.array(img).transpose((2,0,1))
+
+    mask_raw = np.array(mask)[...,None]>0
+    mask = mask.resize((w_t, h_t))
+
+    mask = np.array(mask)
+    mask = (torch.Tensor(mask)>0).float()
+    img = (torch.Tensor(img)).float()
+    img = (img/255-0.5)/0.5
+    img = img[None]
+    mask = mask[None,None]
+
+    with torch.no_grad():
+        generated,_ = model(
+                {'image':img,'mask':mask},
+                mode='inference')
+    generated = torch.clamp(generated, -1, 1)
+    generated = (generated+1)/2*255
+    generated = generated.cpu().numpy().astype(np.uint8)
+    generated = generated[0].transpose((1,2,0))
+    result = generated*mask_raw+img_raw*(1-mask_raw)
+    result = result.astype(np.uint8)
+
+    result = Image.fromarray(result).resize((w_raw, h_raw))
+    result = np.array(result)
+    result = Image.fromarray(result.astype(np.uint8))
+    result.save(f"static/results/{name}")
+    if save_to_input:
+        result.save(f"static/images/{name}")
+
+@app.route('/', methods=['GET', 'POST'])
+def hello(name=None):
+    if 'example' in request.form:
+        filename= request.form['example']
+        image = Image.open(os.path.join(os.path.join(app.config['UPLOAD_FOLDER'], filename)))
+        W, H = image.size
+        return render_template('hello.html', name=name, image_name=filename, image_width=W,
+                image_height=H,list_examples=list_examples)
+    if request.method == 'POST':
+        if 'file' in request.files:
+            file = request.files['file']
+            if file and allowed_file(file.filename):
+                filename = secure_filename(file.filename)
+                image = Image.open(file)
+                W, H = image.size
+                if max(W, H) > max_size:
+                    ratio = float(max_size) / max(W, H)
+                    W = int(W*ratio)
+                    H = int(H*ratio)
+                    image = image.resize((W, H))
+                    filename = "resize_"+filename
+                image.save(os.path.join(os.path.join(app.config['UPLOAD_FOLDER'], filename)))
+                return render_template('hello.html', name=name, image_name=filename, image_width=W,
+                        image_height=H,list_examples=list_examples)
+            else:
+                filename=list_examples[0]
+                image = Image.open(os.path.join(os.path.join(app.config['UPLOAD_FOLDER'], filename)))
+                W, H = image.size
+                return render_template('hello.html', name=name, image_name=filename, image_width=W, image_height=H,
+                        is_alert=True,list_examples=list_examples)
+        if 'mask' in request.form:
+            filename = request.form['imgname']
+            mask_data = request.form['mask']
+            mask_data = mask_data.replace('data:image/png;base64,', '')
+            mask_data = mask_data.replace(' ', '+')
+            mask = base64.b64decode(mask_data)
+            maskname = '.'.join(filename.split('.')[:-1]) + '.png'
+            maskname = maskname.replace("/","_")
+            maskname = "{}_{}".format(random.randint(0, 1000), maskname)
+            with open(os.path.join('static/masks', maskname), "wb") as fh:
+                fh.write(mask)
+            mask = io.BytesIO(mask)
+            mask = Image.open(mask).convert("L")
+            image = Image.open(os.path.join(os.path.join(app.config['UPLOAD_FOLDER'], filename)))
+            W, H = image.size
+            list_op = ["result"]
+            for op in list_op:
+                process_image(image, mask, f"{op}_"+maskname, op, save_to_input=True)
+            return render_template('hello.html', name=name, image_name=filename, #f"{args.opt[0]}_"+maskname
+                    mask_name=maskname, image_width=W, image_height=H, list_opt=list_op,list_examples=list_examples)
+    else:
+        filename=list_examples[0]
+        image = Image.open(os.path.join(os.path.join(app.config['UPLOAD_FOLDER'], filename)))
+        W, H = image.size
+        return render_template('hello.html', name=name, image_name=filename, image_width=W, image_height=H,
+                list_examples=list_examples)
+
+
+
+if __name__ == "__main__":
+
+    app.run(host='0.0.0.0', debug=True, port=port, threaded=True)

demo.sh

@@ -0,0 +1,10 @@
+python demo.py \
+	--name objrmv \
+	--dataset_mode testimage \
+	--model inpaint \
+	--netG baseconv \
+        --which_epoch latest \
+	--image_dir ./datasets/places2sample1k_val/places2samples1k_crop256 \
+	--mask_dir ./datasets/places2sample1k_val/places2samples1k_256_mask_square128 \
+        --output_dir ./results \
+	--port 8897 \

options/test_options.py

@@ -9,6 +9,7 @@
 class TestOptions(BaseOptions):
     def initialize(self, parser):
         BaseOptions.initialize(self, parser)
+        parser.add_argument('--port', type=int, default=8897)
         parser.add_argument('--dataset_mode', type=str, default='coco')
         parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
         parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')

static/images/example.txt

@@ -0,0 +1,19 @@
+images/Places365_val_00020514.png
+images/Places365_val_00016309.png
+images/Places365_val_00007134.png
+images/Places365_val_00015508.png
+images/Places365_val_00017950.png
+images/Places365_val_00007374.png
+images/Places365_val_00012680.png
+images/Places365_val_00014413.png
+images/Places365_val_00032255.png
+images/Places365_val_00015682.png
+images/Places365_val_00014142.png
+images/Places365_val_00026857.png
+images/Places365_val_00004640.png
+images/Places365_val_00031625.png
+images/Places365_val_00001219.png
+images/Places365_val_00026150.png
+images/Places365_val_00005438.png
+images/Places365_val_00030416.png
+images/Places365_val_00028080.png

static/images/images

@@ -0,0 +1 @@
+../../examples/places/images