Add option for keeping dimensions

app.py

@@ -242,12 +242,12 @@ def prepare_model(model_prefix):
     return encoder, decoder
 
 
-def encode(model_prefix, img):
+def encode(model_prefix, img, keep_dims):
     encoder, _ = prepare_model(model_prefix)
-    img_transform = transforms.Compose([
-        transforms.PILToTensor(),
-        transforms.ConvertImageDtype(torch.float)
-    ])
+    img_transform = transforms.Compose(
+        [transforms.PILToTensor(),
+         transforms.ConvertImageDtype(torch.float)] +
+        ([transforms.Resize((224, 224))] if not keep_dims else []))
 
     with torch.no_grad():
         img = img_transform(img.convert("RGB")).unsqueeze(0).to(device)
@@ -276,7 +276,9 @@ def decode(model_prefix, z_str):
 
 
 st.title("Clip Guided Binary Autoencoder")
-st.write("Model is from [@BlinkDL](https://huggingface.co/BlinkDL/clip-guided-binary-autoencoder)")
+st.write(
+    "Model is from [@BlinkDL](https://huggingface.co/BlinkDL/clip-guided-binary-autoencoder)"
+)
 model_prefix = st.selectbox('The model to use',
                             ('out-v7c_d8_256-224-13bit-OB32x0.5-745',
                              'out-v7d_d16_512-224-13bit-OB32x0.5-2487',
@@ -286,13 +288,14 @@ encoder_tab, decoder_tab = st.tabs(["Encode", "Decode"])
 
 with encoder_tab:
     col_in, col_out = st.columns(2)
+    keep_dims = col_in.checkbox('Keep the size of original input image', True)
     uploaded_file = col_in.file_uploader('Choose an Image')
     if uploaded_file is not None:
         image = Image.open(uploaded_file)
         col_in.image(image, 'Input Image')
-        z_str = encode(model_prefix, image)
+        z_str = encode(model_prefix, image, keep_dims)
         col_out.write("Encoded to:")
-        col_out.code(z_str,language=None)
+        col_out.code(z_str, language=None)
         col_out.image(decode(model_prefix, z_str), 'Output Image preview')
 
 with decoder_tab: