Update app.py

app.py

@@ -1,12 +1,14 @@
 import os
- 
-os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116085&Signature=GlUNW6%2B8FxvxWmE9jKIZYOOciKQ%3D" -O weights/RetinaFace-R50.pth')
-os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116208&Signature=hBgvVvKVSNGeXqT8glG%2Bd2t2OKc%3D" -O weights/GPEN-512.pth')
-os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116315&Signature=9tPavW2h%2F1LhIKiXj73sTQoWqcc%3D" -O weights/GPEN-1024-Color.pth ')
-os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x2.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1962694780&Signature=lI%2FolhA%2FyigiTRvoDIVbtMIyhjI%3D" -O weights/realesrnet_x2.pth ')
-os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Inpainting-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116338&Signature=tvYhdLaLgW7UdcUrApXp2jsek8w%3D" -O weights/GPEN-Inpainting-1024.pth ')
-jksp = os.environ['GPEN-BFR-2048']
-os.system(f'wget "{jksp}" -O weights/GPEN-BFR-2048.pth > jksp.log')
+import glob
+
+if len(glob.glob('weights/*.pth')) != 6:
+    os.system('wget -q "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116085&Signature=GlUNW6%2B8FxvxWmE9jKIZYOOciKQ%3D" -O weights/RetinaFace-R50.pth')
+    os.system('wget -q "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116208&Signature=hBgvVvKVSNGeXqT8glG%2Bd2t2OKc%3D" -O weights/GPEN-512.pth')
+    os.system('wget -q "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116315&Signature=9tPavW2h%2F1LhIKiXj73sTQoWqcc%3D" -O weights/GPEN-1024-Color.pth ')
+    os.system('wget -q "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x2.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1962694780&Signature=lI%2FolhA%2FyigiTRvoDIVbtMIyhjI%3D" -O weights/realesrnet_x2.pth ')
+    os.system('wget -q "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Inpainting-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116338&Signature=tvYhdLaLgW7UdcUrApXp2jsek8w%3D" -O weights/GPEN-Inpainting-1024.pth ')
+    jksp = os.environ['GPEN-BFR-2048']
+    os.system(f'wget -q "{jksp}" -O weights/GPEN-BFR-2048.pth')
 
 import gradio as gr
 
@@ -17,65 +19,10 @@ import gradio as gr
 import os
 import cv2
 import glob
-import time
-import math
-import argparse
-import numpy as np
-from PIL import Image, ImageDraw
-import __init_paths
 from face_enhancement import FaceEnhancement 
 from face_colorization import FaceColorization 
 from face_inpainting import FaceInpainting 
-  
-def brush_stroke_mask(img, color=(255,255,255)):
-    min_num_vertex = 8
-    max_num_vertex = 28
-    mean_angle = 2*math.pi / 5
-    angle_range = 2*math.pi / 15
-    min_width = 12
-    max_width = 80
-    def generate_mask(H, W, img=None):
-        average_radius = math.sqrt(H*H+W*W) / 8
-        mask = Image.new('RGB', (W, H), 0)
-        if img is not None: mask = img #Image.fromarray(img)
-
-        for _ in range(np.random.randint(1, 4)):
-            num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
-            angle_min = mean_angle - np.random.uniform(0, angle_range)
-            angle_max = mean_angle + np.random.uniform(0, angle_range)
-            angles = []
-            vertex = []
-            for i in range(num_vertex):
-                if i % 2 == 0:
-                    angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
-                else:
-                    angles.append(np.random.uniform(angle_min, angle_max))
-
-            h, w = mask.size
-            vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
-            for i in range(num_vertex):
-                r = np.clip(
-                    np.random.normal(loc=average_radius, scale=average_radius//2),
-                    0, 2*average_radius)
-                new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
-                new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
-                vertex.append((int(new_x), int(new_y)))
-
-            draw = ImageDraw.Draw(mask)
-            width = int(np.random.uniform(min_width, max_width))
-            draw.line(vertex, fill=color, width=width)
-            for v in vertex:
-                draw.ellipse((v[0] - width//2,
-                              v[1] - width//2,
-                              v[0] + width//2,
-                              v[1] + width//2),
-                             fill=color)
-
-        return mask
 
-    width, height = img.size
-    mask = generate_mask(height, width, img)
-    return mask
 
 def inference(file, mode):
 
@@ -101,7 +48,6 @@ def inference(file, mode):
     elif mode == "inpainting":
         model = {'name':'GPEN-Inpainting-1024', 'size':1024}
         faceinpainter = FaceInpainting(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
-        im = np.asarray(brush_stroke_mask(Image.fromarray(im)))
         inpaint = faceinpainter.process(im)
 
         cv2.imwrite(os.path.join("output.png"), inpaint)
@@ -140,4 +86,4 @@ gr.Interface(
     ['selfie.png', 'selfie']
     ],
     enable_queue=True
-    ).launch()
+).launch()