Update app.py

app.py

@@ -66,35 +66,30 @@ def inferRestoration(img, model_name):
     result = transforms.ToPILImage()(result)
     return result
 
-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
+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)
@@ -129,37 +124,31 @@ def run_cmd(command):
         sys.exit(1)
 
 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)
-
-        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")
+    Restoration_mode == "BOPBTL"
+    if os.path.isdir("Temp"):
         shutil.rmtree("Temp")
     
-    elif Restoration_mode == "Pix2Pix":
-        result_restoration = inferRestoration(image, Restoration_mode)
-    print("Restoration_mode",Restoration_mode)
+    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)
 
+    result_restoration = Image.open("Temp/final_output/input_img.png")
+    shutil.rmtree("Temp")
+    
     result_colorization = inferColorization(result_restoration,Colorizaition_mode)
 
     return result_colorization
  
-examples = [['example/1.jpeg',"BOPBTL","Deoldify"],['example/2.jpg',"BOPBTL","Deoldify"],['example/3.jpg',"BOPBTL","Deoldify"],['example/4.jpg',"BOPBTL","Deoldify"]]
 iface = gr.Interface(run, 
-                    [gr.inputs.Image(),gr.inputs.Radio(["BOPBTL", "Pix2Pix"]),gr.inputs.Radio(["Deoldify", "Pix2Pix Resnet 9block","Pix2Pix Unet 256"])],
-                    outputs="image",
-                    examples=examples).launch(debug=True,share=False)
+                    [gr.inputs.Image(),gr.inputs.Radio(["BOPBTL", "Pix2Pix"]),gr.inputs.Radio(["Deoldify"])],
+                    outputs="image").launch(debug=True,share=False)