Update colorization.py

colorization.py

@@ -6,12 +6,10 @@ from pathlib import Path
 from datetime import datetime
 from typing import Optional
 
-
 model = hub.Module(name='deoldify')
 # NOTE:  Max is 45 with 11GB video cards. 35 is a good default
 render_factor=35
 
-
 def colorize_image(image):
     # now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
     if not os.path.exists("./output"):
@@ -19,13 +17,12 @@ def colorize_image(image):
     # if image is not None:
     #     image.save(f"./output/{now}-input.jpg")    
     model.predict(image.name)
-    
     return f'./output/DeOldify/'+Path(image.name).stem+".png"
 
 # def inference(img, version, scale, weight):
-def inference(img, version, scale):
+def inferenceColorize(img, version, scale):
     # weight /= 100
-    print(img, version, scale)
+    print(img, scale)
     try:
         extension = os.path.splitext(os.path.basename(str(img)))[1]
         img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
@@ -45,6 +42,7 @@ def inference(img, version, scale):
             _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         except RuntimeError as error:
             print('Error', error)
+        
         try:
             if scale != 2:
                 interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
@@ -52,22 +50,25 @@ def inference(img, version, scale):
                 output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
         except Exception as error:
             print('wrong scale input.', error)
+            
         if img_mode == 'RGBA':  # RGBA images should be saved in png format
             extension = 'png'
         else:
             extension = 'jpg'
-        save_path = f'output/out.{extension}'
+        
+        model.predict(img.name)
+        save_path = f'output/DeOldify/'+Path(img.name).stem+'.{extension}'
+    
         cv2.imwrite(save_path, output)
-output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
         return output, save_path
     except Exception as error:
         print('global exception', error)
         return None, None
 
 
-
 colorize = gr.Interface(
-    inference, [
+    inferenceColorize, [
         gr.Markdown("Colorize old black & white photos"),
         gr.Image(type="filepath", label="Input"),