Update app.py

app.py

@@ -1,166 +1,52 @@
 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")
+from PIL import Image 
+import requests
+import random
+r = requests.get(f'https://huggingface.co/spaces/xp3857/bin/raw/main/css.css') 
+css = r.text
+name2 = "xp3857/Image_Restoration_Colorization"
+spaces=[
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+    gr.Interface.load(f"spaces/{name2}"),
+]
+def colorize(input):
+    if input !=None:
+        rn = random.randint(0, 19)
+        space=spaces[rn]        
+        result=space(input)
+        out1 = gr.Pil.update(value=result,visible=True)
+        out2 = gr.Accordion.update(label="Original Image",open=False)
+    else:
+        out1 = None
+        out2 = None
+        pass
+    return out1, out2
+with gr.Blocks(css=css) as myface:
     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)
+        gr.Column()
+        with gr.Column():
+            with gr.Accordion(label="Input Image",open=True) as og:
+                in_win=gr.Pil(label="Input", type="filepath", interactive=True)
+            out_win=gr.Pil(label="Output",visible=False)
+        gr.Column()
+    in_win.change(rem_bg,in_win,[out_win,og])
+myface.queue(concurrency_count=120)
+myface.launch()