Add option for resolution in text box to do worse but faster prediction.

app.py

@@ -16,9 +16,9 @@ device = config.device
 
 
 class ImagePreprocessor():
-    def __init__(self) -> None:
+    def __init__(self, resolution=(1024, 1024)) -> None:
         self.transform_image = transforms.Compose([
-            transforms.Resize((1024, 1024)),
+            transforms.Resize(resolution),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
@@ -42,13 +42,14 @@ model.eval()
 
 # def predict(image_1, image_2):
 #     images = [image_1, image_2]
-def predict(image):
+def predict(image, resolution='1024x1024'):
     images = [image]
     image_shapes = [image.shape[:2] for image in images]
     images = [Image.fromarray(image) for image in images]
 
+    resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
+    image_preprocessor = ImagePreprocessor(resolution=resolution)
     images_proc = []
-    image_preprocessor = ImagePreprocessor()
     for image in images:
         images_proc.append(image_preprocessor.proc(image))
     images_proc = torch.cat([image_proc.unsqueeze(0) for image_proc in images_proc])
@@ -73,6 +74,7 @@ examples = [[_] for _ in glob('materials/examples/*')][:]
 N = 1
 ipt = [gr.Image() for _ in range(N)]
 opt = [gr.Image() for _ in range(N)]
+ipt += [gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution")]
 demo = gr.Interface(
     fn=predict,
     inputs=ipt,
@@ -80,6 +82,6 @@ demo = gr.Interface(
     examples=examples,
     title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',
     description=('Upload a picture, our model will give you the binary maps of the highly accurate segmentation of the salient objects in it. :)'
-                 '\n')
+                 '\nThe resolution used in our training was `1024x1024`, which is too much burden for the huggingface free spaces like this one (cost ~500s). Please set resolution as more than `768x768` for images with many texture details to obtain good results!')
 )
 demo.launch(debug=True)