Update app.py

app.py

@@ -26,22 +26,19 @@ from face_inpainting import FaceInpainting
 
 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)
-    
+    im = cv2.imread(file, cv2.IMREAD_COLOR) if mode != "colorize" else None
     
     if mode == "enhance":
-        return os.path.join("e.png")
+        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")
     elif mode == "colorize":
         model = {'name':'GPEN-1024-Color', 'size':1024}
-        grayf = cv2.imread("e.png", cv2.IMREAD_GRAYSCALE)
+        grayf = cv2.imread(file, 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)
-
         colorf = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
         cv2.imwrite(os.path.join("output.png"), colorf)
         return os.path.join("output.png")
@@ -49,12 +46,10 @@ def inference(file, mode):
         model = {'name':'GPEN-Inpainting-1024', 'size':1024}
         faceinpainter = FaceInpainting(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         inpaint = faceinpainter.process(im)
-
         cv2.imwrite(os.path.join("output.png"), inpaint)
         return os.path.join("output.png")
     elif mode == "selfie":
         model = {'name':'GPEN-BFR-2048', 'size':2048}
-        im = cv2.resize(im, (0,0), fx=4, fy=4)
         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)