Update app.py

app.py

@@ -10,24 +10,15 @@ from skimage import color
 import torchvision.transforms as transforms
 from PIL import Image
 import torch
-import uuid
 import dlib
-uid = uuid.uuid4()
- 
 #os.system("pip install dlib")
 os.system('bash setup.sh')
 
-def run_im(inp):
-    outp=run(inp)
-    return outp
-
-
 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)
     """
@@ -74,30 +65,35 @@ def inferRestoration(img, model_name):
     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
+def inferColorization(img,model_name):
+    #print(model_name)
+    if model_name == "Pix2Pix Resnet 9block":
+        model = torch.hub.load('manhkhanhad/ImageRestorationInfer', 'pix2pixColorization_resnet9b')
+    elif model_name == "Pix2Pix Unet 256":
+        model = torch.hub.load('manhkhanhad/ImageRestorationInfer', 'pix2pixColorization_unet256')
+    elif 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)
@@ -131,42 +127,35 @@ def run_cmd(command):
         print("Process interrupted")
         sys.exit(1)
 
-def run(image):
-    
-    if os.path.isdir("Temp"):
-        shutil.rmtree("Temp")
-    
-    os.makedirs("Temp")
-    os.makedirs("Temp/input")
-    print(type(image))
-    cv2.imwrite("Temp/input/input_img.png", image)
-
-    command = ("python run.py --input_folder "
-            + "Temp/input"
-            + " --output_folder "
-            + "Temp"
-            + " --GPU "
-            + "-1"
-            + " --with_scratch")
-    run_cmd(command)
-    print(f"Temp/final_output: {os.listdir(path='Temp/final_output')}")
-    print(f"Temp: {os.listdir(path='Temp')}")
-    print(f"Temp: {os.listdir(path='Temp/input')}")
-    print(f"Temp: {os.listdir(path='Temp/stage_1_restore_output')}")
-    print(f"Temp: {os.listdir(path='Temp/stage_2_detection_output')}")
-    print(f"Temp: {os.listdir(path='Temp/stage_3_face_output')}")
-    print(f"Temp: {os.listdir(path='Temp/final_output')}")
-    print(f"Save: {os.listdir(path='save')}")
+def run(image,Restoration_mode, Colorizaition_mode):
+    if Restoration_mode == "BOPBTL":
+        if os.path.isdir("Temp"):
+            shutil.rmtree("Temp")
+        
+        os.makedirs("Temp")
+        os.makedirs("Temp/input")
+        print(type(image))
+        cv2.imwrite("Temp/input/input_img.png", image)
 
-    result_restoration = Image.open("Temp/final_output/input_img.png")
-    shutil.rmtree("Temp")
+        command = ("python run.py --input_folder "
+                + "Temp/input"
+                + " --output_folder "
+                + "Temp"
+                + " --GPU "
+                + "-1"
+                + " --with_scratch")
+        run_cmd(command)
 
+        result_restoration = Image.open("Temp/final_output/input_img.png")
+        shutil.rmtree("Temp")
     
-    result_colorization = inferColorization(result_restoration)
+    elif Restoration_mode == "Pix2Pix":
+        result_restoration = inferRestoration(image, Restoration_mode)
+    print("Restoration_mode",Restoration_mode)
+
+    result_colorization = inferColorization(result_restoration,Colorizaition_mode)
 
     return result_colorization
-def load_im(url):
-    return url
  
     
 with gr.Blocks() as app:
@@ -174,8 +163,12 @@ with gr.Blocks() as app:
         gr.Column()
         with gr.Column():
             im = gr.Image(label="Input Image")
+            with gr.Row():
+                rad1 = gr.Radio(["BOPBTL", "Pix2Pix"])
+                rad2 = gr.Radio(["Deoldify", "Pix2Pix Resnet 9block","Pix2Pix Unet 256"])
+                
             im_btn=gr.Button(label="Restore")
             out_im = gr.Image(label="Restored Image")
         gr.Column()
-    im_btn.click(run,im,out_im)
+    im_btn.click(run,[im,rad1,rad2],out_im)
 app.queue(concurrency_count=100).launch(show_api=False)