Add multiple strokes support

app.py

@@ -6,7 +6,7 @@ import numpy as np
 import argparse
 import torch.nn as nn
 import torch.nn.functional as F
-
+import gc
 from baseline.DRL.actor import *
 from baseline.Renderer.stroke_gen import *
 from baseline.Renderer.model import *
@@ -23,6 +23,11 @@ canvas_cnt = divide * divide
 
 Decoder = FCN()
 Decoder.load_state_dict(torch.load(renderer_path))
+actor = ResNet(9, 18, 65) # action_bundle = 5, 65 = 5 * 13
+actor.load_state_dict(torch.load(actor_path))
+actor = actor.to(device).eval()
+Decoder = Decoder.to(device).eval()
+
 
 def decode(x, canvas): # b * (10 + 3)
     x = x.view(-1, 10 + 3)
@@ -37,6 +42,9 @@ def decode(x, canvas): # b * (10 + 3)
     for i in range(5):
         canvas = canvas * (1 - stroke[:, i]) + color_stroke[:, i]
         res.append(canvas)
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache() 
     return canvas, res
 
 def small2large(x):
@@ -86,10 +94,7 @@ def save_img(res, imgid, origin_shape, output_name, divide=False):
     output = cv2.resize(output, origin_shape)
     cv2.imwrite(output_name +"/" + str(imgid) + '.jpg', output)
 
-actor = ResNet(9, 18, 65) # action_bundle = 5, 65 = 5 * 13
-actor.load_state_dict(torch.load(actor_path))
-actor = actor.to(device).eval()
-Decoder = Decoder.to(device).eval()
+
 
 
 
@@ -162,6 +167,41 @@ def paint_img(img):
 
         yield output
 
+
+def change_model(choice: str):
+    global Decoder, actor
+    if choice == "Default":
+        actor_path = 'ckpts/actor.pkl'
+        renderer_path = 'ckpts/renderer.pkl'
+    elif choice == "Triangle":
+        actor_path = 'ckpts/actor_triangle.pkl'
+        renderer_path = 'ckpts/triangle.pkl'  
+    elif choice == "Round":
+        actor_path = 'ckpts/actor_round.pkl'
+        renderer_path = 'ckpts/round.pkl'  
+    else:
+        actor_path = 'ckpts/actor_notrans.pkl'
+        renderer_path = 'ckpts/bezierwotrans.pkl'                
+
+    Decoder.load_state_dict(torch.load(renderer_path, map_location= "cpu"))
+    actor.load_state_dict(torch.load(actor_path, map_location= "cpu"))
+    actor = actor.to(device).eval()
+    Decoder = Decoder.to(device).eval()
+
+
+def wrapper(func):
+    def inner(*args, **kwargs):
+        val = args[0]
+        args_ = tuple(x for i,x in enumerate(args) if i > 0)
+        event = func(*args_, **kwargs)
+        for i in event:
+            if val == "Cancel":
+                yield i
+            else:
+                event.close()
+                break
+    return inner
+
 examples = [
     ["image/chaoyue.png"],
     ["image/degang.png"],
@@ -170,6 +210,28 @@ examples = [
     ["image/mayun.png"],
 ]
 
-demo = gr.Interface(fn=paint_img, inputs=gr.Image(), outputs="image", examples = examples)
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input image")
+            with gr.Row():
+                dropdown = gr.Dropdown(['Default', 'Round', 'Triangle', 'Bezier wo trans'], value= 'Default', label= 'Stroke choice')
+            with gr.Row():
+                with gr.Column():
+                    clr_btn = gr.ClearButton([input_image], variant= "stop")
+                with gr.Column():
+                    translate_btn = gr.Button(value="Paint", variant="primary")
+
+        with gr.Column():
+            output = gr.Image(label="Painting Result")
+    dropdown.select(change_model, dropdown)
+    translate_btn.click(lambda x: gr.Button(value="Cancel", variant="stop") if x == "Paint" else gr.Button(value="Paint", variant="primary"), translate_btn, translate_btn).then(wrapper(paint_img), inputs=[translate_btn, input_image], outputs=output)\
+                 .success(lambda x: gr.Button(value="Paint", variant="primary"), translate_btn, translate_btn)
+    examples = gr.Examples(examples=examples,
+                           inputs=[input_image], cache_examples = False)
+
+
+
+# demo = gr.Interface(fn=paint_img, inputs=gr.Image(), outputs="image", examples = examples)
 demo.queue()
 demo.launch(server_name="0.0.0.0")

ckpts/actor_notrans.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9178218af4256659e4589b378389aa115d2ed293d80d86ac9a650283824b0d6e
+size 44898587

ckpts/actor_round.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29ac4cfddc16642443f3954326055b04a9589d739ae36058ac95c54552990ad7
+size 44898567

ckpts/actor_triangle.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68fefdc1bff96efe8220ed267442b32fd4cc1efdf02701fee0d74315df07997f
+size 44898567

ckpts/bezierwotrans.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82f9e5272d54bc17062f25392e8a0a6070e475852cbcf932ef32818db2cb2fea
+size 44165801

ckpts/round.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a057a399a5f5f481e9606045f133ce79b62c5bc446dc151cb0620b2745f31913
+size 44165813

ckpts/triangle.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59c8fbb8b2ebabf732ea03b9d6904647faa8c074c791df6730b9ddfd1f61cfeb
+size 44165813