Update app.py

app.py

@@ -31,20 +31,20 @@ def lab2rgb(L, AB):
     rgb = color.lab2rgb(Lab) * 255
     return rgb
 
-def get_transform(params=None, grayscale=False, method=Image.BICUBIC):
+def get_transform(model_name,params=None, grayscale=False, method=Image.BICUBIC):
     #params
-    preprocess = 'resize_and_crop'
+    preprocess = 'resize'
     load_size = 256
     crop_size = 256
     transform_list = []
     if grayscale:
         transform_list.append(transforms.Grayscale(1))
-    if 'resize' in preprocess:
+    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))
+    # if 'crop' in preprocess:
+    #     if params is None:
+    #         transform_list.append(transforms.RandomCrop(crop_size))
 
     return transforms.Compose(transform_list)
 
@@ -67,7 +67,7 @@ def inferRestoration(img, model_name):
     return result
 
 def inferColorization(img,model_name):
-    print(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":
@@ -96,10 +96,12 @@ def inferColorization(img,model_name):
         image_pil = transforms.ToPILImage()(result)
         return image_pil
     
-    transform_seq = get_transform()
-    im = transform_seq(img)
-    im = np.array(img)
-    lab = color.rgb2lab(im).astype(np.float32)
+    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
@@ -160,4 +162,4 @@ examples = [['example/1.jpeg',"BOPBTL","Deoldify"],['example/2.jpg',"BOPBTL","De
 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)
+                    examples=examples).launch(debug=True,share=True)