Test

app.py

@@ -1,6 +1,5 @@
 import os
 
-os.system('nvidia-smi')
 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 ')
@@ -20,24 +19,17 @@ from face_colorization import FaceColorization
 
 def inference(file, mode):
 
+    im = cv2.imread(file, cv2.IMREAD_COLOR)
+    im = cv2.resize(im, (0,0), fx=2, fy=2)
+    faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=False)
+    img, orig_faces, enhanced_faces = faceenhancer.process(im)
+    cv2.imwrite(os.path.join("e.png"), img)
+    
     if mode == "enhance":
-        model = {'name':'GPEN-512', 'size':512}
-        im = cv2.imread(file, cv2.IMREAD_COLOR)
-        im = cv2.resize(im, (0,0), fx=2, fy=2)
-        faceenhancer = FaceEnhancement(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
-        img, orig_faces, enhanced_faces = faceenhancer.process(im)
-        cv2.imwrite(os.path.join("output.png"), img)
-        return os.path.join("output.png")
-    else:
-        emodel = {'name':'GPEN-512', 'size':512}
-        im = cv2.imread(file, cv2.IMREAD_COLOR)
-        im = cv2.resize(im, (0,0), fx=2, fy=2)
-        faceenhancer = FaceEnhancement(size=model['size'], model=emodel['name'], channel_multiplier=2, device='cpu')
-        img, orig_faces, enhanced_faces = faceenhancer.process(im)
-        cv2.imwrite(os.path.join("enhanced.png"), img)
-        
+        return os.path.join("e.png")
+    else if mode == "colorize":
         model = {'name':'GPEN-1024-Color', 'size':1024}
-        grayf = cv2.imread("enhanced.png", cv2.IMREAD_GRAYSCALE)
+        grayf = cv2.imread("e.png", cv2.IMREAD_GRAYSCALE)
         grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR) # channel: 1->3
         facecolorizer = FaceColorization(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         colorf = facecolorizer.process(grayf)
@@ -45,6 +37,12 @@ def inference(file, mode):
         colorf = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
         cv2.imwrite(os.path.join("output.png"), colorf)
         return os.path.join("output.png")
+    else:
+        faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=False)
+        img, orig_faces, enhanced_faces = faceenhancer.process(im)
+        cv2.imwrite(os.path.join("output.png"), img)
+        return os.path.join("output.png")
+        
         
 title = "GPEN"
 description = "Gradio demo for GAN Prior Embedded Network for Blind Face Restoration in the Wild. This version of gradio demo includes face colorization from GPEN. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
@@ -54,7 +52,7 @@ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2105.060
 
 gr.Interface(
     inference, 
-    [gr.inputs.Image(type="filepath", label="Input"),gr.inputs.Radio(["enhance","colorize"], type="value", default="enhance", label="Type")], 
+    [gr.inputs.Image(type="filepath", label="Input"),gr.inputs.Radio(["enhance", "colorize", "enhanced+background"], type="value", default="enhance", label="Type")], 
     gr.outputs.Image(type="file", label="Output"),
     title=title,
     description=description,