Update colorization.py

colorization.py

@@ -6,9 +6,11 @@ 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=40
+render_factor=35
+
 
 def colorize_image(image):
     # now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
@@ -17,86 +19,17 @@ 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 inferenceColorize(img, scale):
-    # weight /= 100
-    print(img, scale)
-    '''
-    try:
-        extension = os.path.splitext(os.path.basename(str(img)))[1]
-        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
-        if len(img.shape) == 3 and img.shape[2] == 4:
-            img_mode = 'RGBA'
-        elif len(img.shape) == 2:  # for gray inputs
-            img_mode = None
-            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-        else:
-            img_mode = None
-        h, w = img.shape[0:2]
-        if h < 300:
-            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    return './output/DeOldify/'+Path(image.name).stem+".png"
 
-        try:
-            # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
-            _, _, 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
-                h, w = img.shape[0:2]
-                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'
-    '''
-    os.makedirs('output', exist_ok=True)
-    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)
-    '''
-    return img, save_path
-    '''
-    except Exception as error:
-        print('global exception', error)
-        return None, None
-    '''
 
-colorize = gr.Interface(
-    inferenceColorize, [
-        gr.Markdown("Colorize old black & white photos"),
-        gr.Image(type="filepath", label="Input"),
-        
-        gr.inputs.Number(label="Rescaling factor", default=2),
-        gr.Button("Upscale & Colorize!")
-        
-        # 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 colorized image)"),
-        gr.outputs.File(label="Download the colorized image")
-    ])
-    
-colorize.queue(concurrency_count=4)
-colorize.launch()
-'''
 def create_interface():
     with gr.Blocks() as enhancer:
         gr.Markdown("Colorize old black & white photos")
         with gr.Column(scale=1, label = "Colorize photo", visible=True) as colorize_column:
             colorize_input = gr.Image(type="file")
             colorize_button = gr.Button("Colorize!")
-            colorize_output = gr.outputs.Image(type="numpy", label="Output (The Colorized image)")
-            download_colorize_button = gr.outputs.File(label="Download the colorized image!")
+            colorize_output = gr.outputs.Image(type="numpy", label="Output(The colorized image)")
+            download_colorize_button = gr.outputs.File(label="Download colorized image!")
             colorize_button.click(colorize_image, inputs=colorize_input, outputs=colorize_output)
     enhancer.launch()
 
@@ -104,5 +37,4 @@ def run_code():
     create_interface()
 
 # The main function
-run_code()
-'''
+run_code()