Create app.py

app.py

@@ -0,0 +1,166 @@
+import gradio as gr
+import os
+import cv2 
+import shutil
+import sys
+from subprocess import call
+import torch
+import numpy as np
+from skimage import color
+import torchvision.transforms as transforms
+from PIL import Image
+import torch
+import uuid
+ 
+#os.system("pip install dlib")
+os.system('bash setup.sh')
+
+def lab2rgb(L, AB):
+    """Convert an Lab tensor image to a RGB numpy output
+    Parameters:
+        L  (1-channel tensor array): L channel images (range: [-1, 1], torch tensor array)
+        AB (2-channel tensor array):  ab channel images (range: [-1, 1], torch tensor array)
+    Returns:
+        rgb (RGB numpy image): rgb output images  (range: [0, 255], numpy array)
+    """
+    AB2 = AB * 110.0
+    L2 = (L + 1.0) * 50.0
+    Lab = torch.cat([L2, AB2], dim=1)
+    Lab = Lab[0].data.cpu().float().numpy()
+    Lab = np.transpose(Lab.astype(np.float64), (1, 2, 0))
+    rgb = color.lab2rgb(Lab) * 255
+    return rgb
+
+def get_transform(model_name,params=None, grayscale=False, method=Image.BICUBIC):
+    #params
+    preprocess = 'resize'
+    load_size = 256
+    crop_size = 256
+    transform_list = []
+    if grayscale:
+        transform_list.append(transforms.Grayscale(1))
+    if model_name == "Pix2Pix Unet 256":
+        osize = [load_size, load_size]
+        transform_list.append(transforms.Resize(osize, method))
+    # if 'crop' in preprocess:
+    #     if params is None:
+    #         transform_list.append(transforms.RandomCrop(crop_size))
+
+    return transforms.Compose(transform_list)
+
+def inferRestoration(img, model_name):
+    #if model_name == "Pix2Pix":
+    model = torch.hub.load('manhkhanhad/ImageRestorationInfer', 'pix2pixRestoration_unet256')
+    transform_list = [
+                transforms.ToTensor(),
+                transforms.Resize([256,256], Image.BICUBIC),
+                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+                ]
+    transform = transforms.Compose(transform_list)
+    img = transform(img)
+    img = torch.unsqueeze(img, 0)
+    result = model(img)
+    result = result[0].detach()
+    result = (result +1)/2.0
+    
+    result = transforms.ToPILImage()(result)
+    return result
+
+def inferColorization(img):
+    model_name = "Deoldify"
+    model = torch.hub.load('manhkhanhad/ImageRestorationInfer', 'DeOldifyColorization')
+    transform_list = [
+              transforms.ToTensor(),
+              transforms.Normalize((0.5,), (0.5,))
+                ]
+    transform = transforms.Compose(transform_list)
+    #a = transforms.ToTensor()(a)
+    img = img.convert('L')
+    img = transform(img)
+    img = torch.unsqueeze(img, 0)
+    result = model(img)
+    
+    result = result[0].detach()
+    result = (result +1)/2.0
+    
+    #img = transforms.Grayscale(3)(img)
+    #img = transforms.ToTensor()(img)
+    #img = torch.unsqueeze(img, 0)
+    #result = model(img)
+    #result = torch.clip(result, min=0, max=1)
+    image_pil = transforms.ToPILImage()(result)
+    return image_pil
+    
+    transform_seq = get_transform(model_name)
+    img = transform_seq(img)
+    # if model_name == "Pix2Pix Unet 256":
+    #     img.resize((256,256))
+    img = np.array(img)
+    lab = color.rgb2lab(img).astype(np.float32)
+    lab_t = transforms.ToTensor()(lab)
+    A = lab_t[[0], ...] / 50.0 - 1.0
+    B = lab_t[[1, 2], ...] / 110.0
+    #data = {'A': A, 'B': B, 'A_paths': "", 'B_paths': ""}
+    L = torch.unsqueeze(A, 0)
+    #print(L.shape)
+    ab = model(L)
+    Lab = lab2rgb(L, ab).astype(np.uint8)
+    image_pil = Image.fromarray(Lab)
+    #image_pil.save('test.png')
+    #print(Lab.shape)
+    return image_pil
+    
+def colorizaition(image,model_name):
+    image = Image.fromarray(image)
+    result = inferColorization(image,model_name)
+    return result
+
+
+def run_cmd(command):
+    try:
+        call(command, shell=True)
+    except KeyboardInterrupt:
+        print("Process interrupted")
+        sys.exit(1)
+
+def run(image):
+    uid = uuid.uuid4()
+    
+    if os.path.isdir(f"Temp{uid}"):
+        shutil.rmtree(f"Temp{uid}")
+    
+    os.makedirs(f"Temp{uid}")
+    os.makedirs(f"Temp{uid}/input")
+    print(type(image))
+    cv2.imwrite(f"Temp{uid}/input/input_img.png", image)
+
+    command = ("python run.py --input_folder "
+            + f"Temp{uid}/input"
+            + " --output_folder "
+            + f"Temp{uid}"
+            + " --GPU "
+            + "-1"
+            + " --with_scratch")
+    run_cmd(command)
+
+    result_restoration = Image.open(f"Temp{uid}/final_output/input_img.png")
+    shutil.rmtree(f"Temp{uid}")
+    
+    result_colorization = inferColorization(result_restoration)
+
+    return result_colorization
+def load_im(url):
+    return url
+ 
+    
+with gr.Blocks() as app:
+    im = gr.Image(label="Input Image")
+    with gr.Row():
+        im_u = gr.Textbox()
+        lim_btn=gr.Button("Load")
+    im_btn=gr.Button(label="Restore")
+    out_im = gr.Image(label="Restored Image")
+
+    #lim_btn(load_im,im_u,im)
+    im_btn.click(run,[im,im_u],out_im)
+app.queue(concurrency_count=100).launch(show_api=False)