Update app.py

app.py

@@ -11,8 +11,8 @@ import torchvision.transforms as transforms
 from PIL import Image
 import torch
 
-os.system("pip install dlib")
-os.system('bash setup.sh')
+#os.system("pip install dlib")
+#os.system('bash setup.sh')
 
 def lab2rgb(L, AB):
     """Convert an Lab tensor image to a RGB numpy output
@@ -48,11 +48,29 @@ def get_transform(params=None, grayscale=False, method=Image.BICUBIC):
 
     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):
     print(model_name)
-    if model_name == "Pix2Pix_resnet9b":
+    if model_name == "Pix2Pix Resnet 9block":
         model = torch.hub.load('manhkhanhad/ImageRestorationInfer', 'pix2pixColorization_resnet9b')
-    elif model_name == "Pix2Pix_unet256":
+    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')
@@ -108,30 +126,37 @@ 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)
+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)
+        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_restoration = Image.open("Temp/final_output/input_img.png")
+        shutil.rmtree("Temp")
+    
+    elif Restoration_mode == "Pix2Pix":
+        result_restoration = inferRestoration(image, Restoration_mode)
 
-    result_colorization = inferColorization(result_restoration,"Deoldify")
+    result_colorization = inferColorization(result_restoration,colorizaition_mode)
 
     return result_colorization
  
 examples = [['example/1.jpeg'],['example/2.jpg'],['example/3.jpg'],['example/4.jpg']]
-iface = gr.Interface(fn=run, inputs="image", outputs="image",examples=examples).launch(debug=True,share=False)
+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=True)