Update app.py

app.py

@@ -22,18 +22,19 @@ if not os.path.exists('RestoreFormer.pth'):
 if not os.path.exists('CodeFormer.pth'):
     os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth -P .")
 
-torch.hub.download_url_to_file(
-    'https://haoluobo.com/wp-content/uploads/2023/01/%E6%A8%B1%E6%9C%A8%E8%8A%B1%E9%81%93.jpg',
-    'HanamichiSakuragi.jpg')
-torch.hub.download_url_to_file(
-    'https://haoluobo.com/wp-content/uploads/2023/01/%E6%9D%8E%E4%B8%96%E6%B0%91.jpg',
-    'LiShiming.jpg')
-torch.hub.download_url_to_file(
-    'https://haoluobo.com/wp-content/uploads/2023/01/%E4%B9%BE%E9%9A%86.jpg',
-    'QianLong.jpg')
-torch.hub.download_url_to_file(
-    'https://user-images.githubusercontent.com/17445847/187401133-8a3bf269-5b4d-4432-b2f0-6d26ee1d3307.png',
-    '10045.png')
+if not os.path.exists('HanamichiSakuragi.jpg'):
+    torch.hub.download_url_to_file(
+        'https://haoluobo.com/wp-content/uploads/2023/01/%E6%A8%B1%E6%9C%A8%E8%8A%B1%E9%81%93.jpg',
+        'HanamichiSakuragi.jpg')
+    torch.hub.download_url_to_file(
+        'https://haoluobo.com/wp-content/uploads/2023/01/%E6%9D%8E%E4%B8%96%E6%B0%91.jpg',
+        'LiShiming.jpg')
+    torch.hub.download_url_to_file(
+        'https://haoluobo.com/wp-content/uploads/2023/01/%E4%B9%BE%E9%9A%86.jpg',
+        'QianLong.jpg')
+    torch.hub.download_url_to_file(
+        'https://user-images.githubusercontent.com/17445847/187401133-8a3bf269-5b4d-4432-b2f0-6d26ee1d3307.png',
+        '10045.png')
 
 # background enhancer with RealESRGAN
 model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
@@ -90,9 +91,13 @@ def inference(img, version, scale, blur_face):
             # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
             face_helper = face_enhancer.face_helper
             align_warp_face = face_helper.align_warp_face
+
             def new_align_warp_face(*args, **kwargs):
                 align_warp_face(*args, **kwargs)  # save_cropped_path
+                face_helper.org_cropped_faces = face_helper.cropped_faces
                 face_helper.cropped_faces = [cv2.GaussianBlur(e, (blur_face, blur_face), 0) for e in face_helper.cropped_faces]
+                print("find face count:", len(face_helper.cropped_faces))
+
             face_helper.align_warp_face = new_align_warp_face
             _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         except RuntimeError as error:
@@ -113,7 +118,7 @@ def inference(img, version, scale, blur_face):
         cv2.imwrite(save_path, output)
 
         output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-        return output, save_path
+        return output, save_path, [cv2.cvtColor(e, cv2.COLOR_BGR2RGB) for e in face_enhancer.face_helper.org_cropped_faces]
     except Exception as error:
         print('global exception', error)
         return None, None
@@ -146,7 +151,8 @@ demo = gr.Interface(
         # gr.Slider(0, 100, label='Weight, only for CodeFormer. 0 for better quality, 100 for better identity', default=50)
     ], [
         gr.outputs.Image(type="numpy", label="Output (The whole image)"),
-        gr.outputs.File(label="Download the output image")
+        gr.outputs.File(label="Download the output image"),
+        gr.Gallery(label="All face").style(grid=[2], height="auto")
     ],
     title=title,
     description=description,