Update app.py

app.py

@@ -11,18 +11,17 @@ from torchvision import transforms
 import torchtext
 
 
-# Images
-torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2021/08/04/14/16/tower-6521842_1280.jpg', 'tower.jpg')
-torch.hub.download_url_to_file('https://cdn.pixabay.com/photo/2017/08/31/05/36/buildings-2699520_1280.jpg', 'city.jpg')
 
 idx = 0
 
 torchtext.utils.download_from_url("https://drive.google.com/uc?id=1NDD54BLligyr8tzo8QGI5eihZisXK1nq", root=".")
 
 
-def save_img(img, output_path):
+def to_PIL_img(img):
     result = Image.fromarray((img.data.cpu().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8))
-    result.save(output_path)
+    return result
+def save_img(img, output_path):
+    to_PIL_img(img).save(output_path)
 
 
 def param2stroke(param, H, W, meta_brushes):
@@ -87,7 +86,7 @@ def param2stroke(param, H, W, meta_brushes):
 
 
 def param2img_serial(
-        param, decision, meta_brushes, cur_canvas, frame_dir, has_border=False, original_h=None, original_w=None):
+        param, decision, meta_brushes, cur_canvas, frame_dir, has_border=False, original_h=None, original_w=None, *, all_frames):
     """
     Input stroke parameters and decisions for each patch, meta brushes, current canvas, frame directory,
     and whether there is a border (if intermediate painting results are required).
@@ -146,7 +145,8 @@ def param2img_serial(
                                            device=this_canvas.device)
         selected_alphas = torch.zeros(selected_param.shape[0], 3, patch_size_y, patch_size_x, device=this_canvas.device)
         if selected_param[selected_decision, :].shape[0] > 0:
-            selected_foregrounds[selected_decision, :, :, :], selected_alphas[selected_decision, :, :, :] =                               param2stroke(selected_param[selected_decision, :], patch_size_y, patch_size_x, meta_brushes)
+            selected_foregrounds[selected_decision, :, :, :], selected_alphas[selected_decision, :, :, :] = \
+                param2stroke(selected_param[selected_decision, :], patch_size_y, patch_size_x, meta_brushes)
         selected_foregrounds = selected_foregrounds.view(
             b, selected_h, selected_w, 3, patch_size_y, patch_size_x).contiguous()
         selected_alphas = selected_alphas.view(b, selected_h, selected_w, 3, patch_size_y, patch_size_x).contiguous()
@@ -164,6 +164,11 @@ def param2img_serial(
         factor = 2
     else:
         factor = 4
+
+    def store_frame(img):
+        all_frames.append(to_PIL_img(img))
+        
+
     if even_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
         for i in range(s):
             canvas = partial_render(cur_canvas, even_y_even_x_coord_y, even_y_even_x_coord_x, i)
@@ -177,6 +182,7 @@ def param2img_serial(
                 frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
                              patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
                 save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+                store_frame(frame[0])
 
     if odd_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
         for i in range(s):
@@ -193,6 +199,7 @@ def param2img_serial(
                 frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
                              patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
                 save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+                store_frame(frame[0])
 
     if odd_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
         for i in range(s):
@@ -208,6 +215,7 @@ def param2img_serial(
                 frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
                              patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
                 save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+                store_frame(frame[0])
 
     if even_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
         for i in range(s):
@@ -223,6 +231,7 @@ def param2img_serial(
                 frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
                              patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
                 save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+                store_frame(frame[0])
 
     cur_canvas = cur_canvas[:, :, patch_size_y // 4:-patch_size_y // 4, patch_size_x // 4:-patch_size_x // 4]
 
@@ -383,10 +392,10 @@ def crop(img, h, w):
 
 
 def main(input_path, model_path, output_dir, need_animation=False, resize_h=None, resize_w=None, serial=False):
-    # if not os.path.exists(output_dir):
-    #     os.mkdir(output_dir)
+    if not os.path.exists(output_dir):
+        os.mkdir(output_dir)
     input_name = os.path.basename(input_path)
-    # output_path = os.path.join(output_dir, input_name)
+    output_path = os.path.join(output_dir, input_name)
     frame_dir = None
     if need_animation:
         if not serial:
@@ -416,6 +425,7 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
         original_img_pad_size = patch_size * (2 ** K)
         original_img_pad = pad(original_img, original_img_pad_size, original_img_pad_size)
         final_result = torch.zeros_like(original_img_pad).to(device)
+        all_frames = []
         for layer in range(0, K + 1):
             layer_size = patch_size * (2 ** layer)
             img = F.interpolate(original_img_pad, (layer_size, layer_size))
@@ -449,7 +459,7 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
             param[..., 2:4] = param[..., 2:4] / 2
             if serial:
                 final_result = param2img_serial(param, decision, meta_brushes, final_result,
-                                                frame_dir, False, original_h, original_w)
+                                                frame_dir, False, original_h, original_w, all_frames = all_frames)
             else:
                 final_result = param2img_parallel(param, decision, meta_brushes, final_result)
 
@@ -486,39 +496,39 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
         param[..., 2:4] = param[..., 2:4] / 2
         if serial:
             final_result = param2img_serial(param, decision, meta_brushes, final_result,
-                                            frame_dir, True, original_h, original_w)
+                                            frame_dir, True, original_h, original_w, all_frames = all_frames)
         else:
             final_result = param2img_parallel(param, decision, meta_brushes, final_result)
         final_result = final_result[:, :, border_size:-border_size, border_size:-border_size]
 
         final_result = crop(final_result, original_h, original_w)
-        # save_img(final_result[0], output_path)
+        save_img(final_result[0], output_path)
         tensor_to_pil = transforms.ToPILImage()(final_result[0].squeeze_(0))
-        return tensor_to_pil
+        #return tensor_to_pil
+
+        all_frames[0].save(os.path.join(frame_dir, 'animation.gif'),
+               save_all=True, append_images=all_frames[1:], optimize=False, duration=40, loop=0)
+        return os.path.join(frame_dir, "animation.gif")
 
 
 def gradio_inference(image):
     return main(input_path=image.name,
          model_path='model.pth',
          output_dir='output/',
-         need_animation=False,  # whether need intermediate results for animation.
+         need_animation=True,  # whether need intermediate results for animation.
          resize_h=512,         # resize original input to this size. None means do not resize.
          resize_w=512,         # resize original input to this size. None means do not resize.
-         serial=False)          # if need animation, serial must be True.
+         serial=True)          # if need animation, serial must be True.
 
-title = "Paint Transformer"
-description = "Gradio demo for Paint Transformer: Feed Forward Neural Painting with Stroke Prediction. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.03798'>Paint Transformer: Feed Forward Neural Painting with Stroke Prediction</a> | <a href='https://github.com/Huage001/PaintTransformer'>Github Repo</a></p>"
+title = "Anime2Sketch"
+description = "demo for Anime2Sketch. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2104.05703'>Adversarial Open Domain Adaption for Sketch-to-Photo Synthesis</a> | <a href='https://github.com/Mukosame/Anime2Sketch'>Github Repo</a></p>"
 
 gr.Interface(
     gradio_inference, 
     [gr.inputs.Image(type="file", label="Input")], 
-    gr.outputs.Image(type="pil", label="Output"),
+    gr.outputs.Image(type="file", label="Output"),
     title=title,
     description=description,
-    article=article,
-    examples=[
-    ['city.jpg'],
-    ['tower.jpg']
-    ]
+    article=article
     ).launch(debug=True)