Update app.py

app.py

@@ -9,25 +9,20 @@ import math
 import gradio as gr
 from torchvision import transforms
 import torchtext
-
+from stat import ST_CTIME
+from datetime import datetime, timedelta
+import shutil
 torch.hub.download_url_to_file('https://i.imgur.com/WEHmKef.jpg', 'gpu.jpg')
-
 # 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 to_PIL_img(img):
     result = Image.fromarray((img.data.cpu().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8))
     return result
 def save_img(img, output_path):
     to_PIL_img(img).save(output_path)
-
-
 def param2stroke(param, H, W, meta_brushes):
     """
     Input a set of stroke parameters and output its corresponding foregrounds and alpha maps.
@@ -38,7 +33,6 @@ def param2stroke(param, H, W, meta_brushes):
         W: output width.
         meta_brushes: a tensor with shape 2 x 3 x meta_brush_height x meta_brush_width.
          The first slice on the batch dimension denotes vertical brush and the second one denotes horizontal brush.
-
     Returns:
         foregrounds: a tensor with shape n_strokes x 3 x H x W, containing color information.
         alphas: a tensor with shape n_strokes x 3 x H x W,
@@ -61,7 +55,6 @@ def param2stroke(param, H, W, meta_brushes):
     index[h > w] = 0
     index[h <= w] = 1
     brush = meta_brushes_resize[index.long()]
-
     # Calculate warp matrix according to the rules defined by pytorch, in order for warping.
     warp_00 = cos_theta / w
     warp_01 = sin_theta * H / (W * w)
@@ -87,8 +80,6 @@ def param2stroke(param, H, W, meta_brushes):
     foreground = morphology.dilation(foreground)
     alphas = morphology.erosion(alphas)
     return foreground, alphas
-
-
 def param2img_serial(
         param, decision, meta_brushes, cur_canvas, frame_dir, has_border=False, original_h=None, original_w=None, *, all_frames):
     """
@@ -111,7 +102,6 @@ def param2img_serial(
          on the border before saving, or there would be a black border.
         original_h: to indicate the original height for cropping when saving intermediate results.
         original_w: to indicate the original width for cropping when saving intermediate results.
-
     Returns:
         cur_canvas: a tensor with shape batch size x 3 x H x W, denoting painting results.
     """
@@ -133,7 +123,6 @@ def param2img_serial(
     odd_y_even_x_coord_y, odd_y_even_x_coord_x = torch.meshgrid([odd_idx_y, even_idx_x])
     cur_canvas = F.pad(cur_canvas, [patch_size_x // 4, patch_size_x // 4,
                                     patch_size_y // 4, patch_size_y // 4, 0, 0, 0, 0])
-
     def partial_render(this_canvas, patch_coord_y, patch_coord_x, stroke_id):
         canvas_patch = F.unfold(this_canvas, (patch_size_y, patch_size_x),
                                 stride=(patch_size_y // 2, patch_size_x // 2))
@@ -161,17 +150,14 @@ def param2img_serial(
         this_canvas = this_canvas.view(b, 3, selected_h * patch_size_y, selected_w * patch_size_x).contiguous()
         # this_canvas: b, 3, selected_h * py, selected_w * px
         return this_canvas
-
     global idx
     if has_border:
         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)
@@ -186,7 +172,6 @@ def param2img_serial(
                              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):
             canvas = partial_render(cur_canvas, odd_y_odd_x_coord_y, odd_y_odd_x_coord_x, i)
@@ -203,7 +188,6 @@ def param2img_serial(
                              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):
             canvas = partial_render(cur_canvas, odd_y_even_x_coord_y, odd_y_even_x_coord_x, i)
@@ -219,7 +203,6 @@ def param2img_serial(
                              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):
             canvas = partial_render(cur_canvas, even_y_odd_x_coord_y, even_y_odd_x_coord_x, i)
@@ -235,12 +218,8 @@ def param2img_serial(
                              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]
-
     return cur_canvas
-
-
 def param2img_parallel(param, decision, meta_brushes, cur_canvas):
     """
         Input stroke parameters and decisions for each patch, meta brushes, current canvas, frame directory,
@@ -255,7 +234,6 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
             The first slice on the batch dimension denotes vertical brush and the second one denotes horizontal brush.
             cur_canvas: a tensor with shape batch size x 3 x H x W,
              where H and W denote height and width of padded results of original images.
-
         Returns:
             cur_canvas: a tensor with shape batch size x 3 x H x W, denoting painting results.
         """
@@ -289,11 +267,8 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
     alphas = alphas.view(-1, h, w, s, 3, patch_size_y, patch_size_x).contiguous()
     # foreground, alpha: b, h, w, stroke_per_patch, 3, render_size_y, render_size_x
     decision = decision.view(-1, h, w, s, 1, 1, 1).contiguous()
-
     # decision: b, h, w, stroke_per_patch, 1, 1, 1
-
     def partial_render(this_canvas, patch_coord_y, patch_coord_x):
-
         canvas_patch = F.unfold(this_canvas, (patch_size_y, patch_size_x),
                                 stride=(patch_size_y // 2, patch_size_x // 2))
         # canvas_patch: b, 3 * py * px, h * w
@@ -317,7 +292,6 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
         this_canvas = this_canvas.view(b, 3, h_half * patch_size_y, w_half * patch_size_x).contiguous()
         # this_canvas: b, 3, h_half * py, w_half * px
         return this_canvas
-
     if even_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
         canvas = partial_render(cur_canvas, even_y_even_x_coord_y, even_y_even_x_coord_x)
         if not is_odd_y:
@@ -325,7 +299,6 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
         if not is_odd_x:
             canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
         cur_canvas = canvas
-
     if odd_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
         canvas = partial_render(cur_canvas, odd_y_odd_x_coord_y, odd_y_odd_x_coord_x)
         canvas = torch.cat([cur_canvas[:, :, :patch_size_y // 2, -canvas.shape[3]:], canvas], dim=2)
@@ -335,7 +308,6 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
         if is_odd_x:
             canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
         cur_canvas = canvas
-
     if odd_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
         canvas = partial_render(cur_canvas, odd_y_even_x_coord_y, odd_y_even_x_coord_x)
         canvas = torch.cat([cur_canvas[:, :, :patch_size_y // 2, :canvas.shape[3]], canvas], dim=2)
@@ -344,7 +316,6 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
         if not is_odd_x:
             canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
         cur_canvas = canvas
-
     if even_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
         canvas = partial_render(cur_canvas, even_y_odd_x_coord_y, even_y_odd_x_coord_x)
         canvas = torch.cat([cur_canvas[:, :, :canvas.shape[2], :patch_size_x // 2], canvas], dim=3)
@@ -353,12 +324,8 @@ def param2img_parallel(param, decision, meta_brushes, cur_canvas):
         if is_odd_x:
             canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
         cur_canvas = canvas
-
     cur_canvas = cur_canvas[:, :, patch_size_y // 4:-patch_size_y // 4, patch_size_x // 4:-patch_size_x // 4]
-
     return cur_canvas
-
-
 def read_img(img_path, img_type='RGB', h=None, w=None):
     img = Image.open(img_path).convert(img_type)
     if h is not None and w is not None:
@@ -369,8 +336,6 @@ def read_img(img_path, img_type='RGB', h=None, w=None):
     img = img.transpose((2, 0, 1))
     img = torch.from_numpy(img).unsqueeze(0).float() / 255.
     return img
-
-
 def pad(img, H, W):
     b, c, h, w = img.shape
     pad_h = (H - h) // 2
@@ -382,8 +347,6 @@ def pad(img, H, W):
     img = torch.cat([torch.zeros((b, c, H, pad_w), device=img.device), img,
                      torch.zeros((b, c, H, pad_w + remainder_w), device=img.device)], dim=-1)
     return img
-
-
 def crop(img, h, w):
     H, W = img.shape[-2:]
     pad_h = (H - h) // 2
@@ -392,11 +355,21 @@ def crop(img, h, w):
     remainder_w = (W - w) % 2
     img = img[:, :, pad_h:H - pad_h - remainder_h, pad_w:W - pad_w - remainder_w]
     return img
-
-
 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)
+
+    for entry in os.listdir(output_dir):
+        path = os.path.join(output_dir, entry)
+        stats = os.stat(path)
+        created_time = datetime.fromtimestamp(stats[ST_CTIME])
+        if created_time < datetime.now() - timedelta(minutes = 10):
+            if os.path.isdir(path):
+                shutil.rmtree(path)
+            else:
+                os.remove(path)
+        
+    
     input_name = os.path.basename(input_path)
     output_path = os.path.join(output_dir, input_name)
     frame_dir = None
@@ -415,12 +388,10 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
     net_g.eval()
     for param in net_g.parameters():
         param.requires_grad = False
-
     brush_large_vertical = read_img('brush/brush_large_vertical.png', 'L').to(device)
     brush_large_horizontal = read_img('brush/brush_large_horizontal.png', 'L').to(device)
     meta_brushes = torch.cat(
         [brush_large_vertical, brush_large_horizontal], dim=0)
-
     with torch.no_grad():
         original_img = read_img(input_path, 'RGB', resize_h, resize_w).to(device)
         original_h, original_w = original_img.shape[-2:]
@@ -438,14 +409,12 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
                                     stride=(patch_size, patch_size))
             # There are patch_num * patch_num patches in total
             patch_num = (layer_size - patch_size) // patch_size + 1
-
             # img_patch, result_patch: b, 3 * output_size * output_size, h * w
             img_patch = img_patch.permute(0, 2, 1).contiguous().view(-1, 3, patch_size, patch_size).contiguous()
             result_patch = result_patch.permute(0, 2, 1).contiguous().view(
                 -1, 3, patch_size, patch_size).contiguous()
             shape_param, stroke_decision = net_g(img_patch, result_patch)
             stroke_decision = network.SignWithSigmoidGrad.apply(stroke_decision)
-
             grid = shape_param[:, :, :2].view(img_patch.shape[0] * stroke_num, 1, 1, 2).contiguous()
             img_temp = img_patch.unsqueeze(1).contiguous().repeat(1, stroke_num, 1, 1, 1).view(
                 img_patch.shape[0] * stroke_num, 3, patch_size, patch_size).contiguous()
@@ -465,7 +434,6 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
                                                 frame_dir, False, original_h, original_w, all_frames = all_frames)
             else:
                 final_result = param2img_parallel(param, decision, meta_brushes, final_result)
-
         border_size = original_img_pad_size // (2 * patch_num)
         img = F.interpolate(original_img_pad, (patch_size * (2 ** layer), patch_size * (2 ** layer)))
         result = F.interpolate(final_result, (patch_size * (2 ** layer), patch_size * (2 ** layer)))
@@ -482,7 +450,6 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
         img_patch = img_patch.permute(0, 2, 1).contiguous().view(-1, 3, patch_size, patch_size).contiguous()
         result_patch = result_patch.permute(0, 2, 1).contiguous().view(-1, 3, patch_size, patch_size).contiguous()
         shape_param, stroke_decision = net_g(img_patch, result_patch)
-
         grid = shape_param[:, :, :2].view(img_patch.shape[0] * stroke_num, 1, 1, 2).contiguous()
         img_temp = img_patch.unsqueeze(1).contiguous().repeat(1, stroke_num, 1, 1, 1).view(
             img_patch.shape[0] * stroke_num, 3, patch_size, patch_size).contiguous()
@@ -503,17 +470,13 @@ def main(input_path, model_path, output_dir, need_animation=False, resize_h=None
         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)
         tensor_to_pil = transforms.ToPILImage()(final_result[0].squeeze_(0))
         #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"), tensor_to_pil
-
-
     
 def gradio_inference(image):
     return main(input_path=image.name,
@@ -523,7 +486,6 @@ def gradio_inference(image):
          resize_h=400,         # resize original input to this size. None means do not resize.
          resize_w=400,         # resize original input to this size. None means do not resize.
          serial=True)          # if need animation, serial must be True.
-
 inferences_running = 0
 def throttled_inference(image):
     global inferences_running
@@ -538,11 +500,9 @@ def throttled_inference(image):
     finally:
         print("Inference finished")
         inferences_running -= 1
-
 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>"
-
 gr.Interface(
     throttled_inference, 
     gr.inputs.Image(type="file", label="Input"), 
@@ -555,4 +515,4 @@ gr.Interface(
     ['city.jpg'],
     ['tower.jpg']
     ]
-    ).launch(debug=True)
+    ).launch(debug=True)