Update app.py

app.py

@@ -19,34 +19,56 @@ import torch
 from basicsr.utils import imwrite
 from gfpgan import GFPGANer
 
+# background upsampler
+if args.bg_upsampler == 'realesrgan':
+    if not torch.cuda.is_available():  # CPU
+        import warnings
+        warnings.warn('The unoptimized RealESRGAN is very slow on CPU. We do not use it. '
+                      'If you really want to use it, please modify the corresponding codes.')
+        bg_upsampler = None
+    else:
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        bg_upsampler = RealESRGANer(
+            scale=2,
+            model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
+            model=model,
+            tile=args.bg_tile,
+            tile_pad=10,
+            pre_pad=0,
+            half=True)  # need to set False in CPU mode
+else:
+    bg_upsampler = None
+    
+# set up GFPGAN restorer
+restorer = GFPGANer(
+    model_path=args.model_path,
+    upscale=args.upscale,
+    arch=args.arch,
+    channel_multiplier=args.channel,
+    bg_upsampler=bg_upsampler)
+img_list = sorted(glob.glob(os.path.join(args.test_path, '*')))
+
+
+
 os.system("pip install gfpgan")
 os.system("pip freeze")
 os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth -P .")
 
 
-def run_cmd(command):
-    try:
-        print(command)
-        call(command, shell=True)
-    except KeyboardInterrupt:
-        print("Process interrupted")
-        sys.exit(1)
-
 def inference(img):
-    _id = randint(1, 10000)
-    INPUT_DIR = "/tmp/input_image" + str(_id) + "/"
-    OUTPUT_DIR = "/tmp/output_image" + str(_id) + "/"
-    run_cmd("rm -rf " + INPUT_DIR)
-    run_cmd("rm -rf " + OUTPUT_DIR)
-    run_cmd("mkdir " + INPUT_DIR)
-    run_cmd("mkdir " + OUTPUT_DIR)
-    basewidth = 256
-    wpercent = (basewidth/float(img.size[0]))
-    hsize = int((float(img.size[1])*float(wpercent)))
-    img = img.resize((basewidth,hsize), Image.ANTIALIAS)
-    img.save(INPUT_DIR + "1.jpg", "JPEG")
-    run_cmd("python inference_gfpgan.py --upscale 2 --test_path "+ INPUT_DIR + " --save_root " + OUTPUT_DIR + " --paste_back")
-    return os.path.join(OUTPUT_DIR, "1_00.png")
+    input_img = cv2.imread(img, cv2.IMREAD_COLOR)
+    # restore faces and background if necessary
+    cropped_faces, restored_faces, restored_img = restorer.enhance(
+        input_img, has_aligned=args.aligned, only_center_face=args.only_center_face, paste_back=args.paste_back)
+    # save faces
+    # save cropped face
+   
+    # save comparison image
+    cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
+   
+    return Image.fromarray(cmp_img)
         
 title = "GFP-GAN"
 description = "Gradio demo for GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Please click submit only once"