fix input arg name

Embed_Secret.py

@@ -91,8 +91,8 @@ def load_UNet(args):
     # config_file = '/mnt/fast/nobackup/scratch4weeks/tb0035/projects/diffsteg/FLAE/simple_t2_croprs/configs/-project.yaml'
     # weight_file = '/mnt/fast/nobackup/scratch4weeks/tb0035/projects/diffsteg/FLAE/simple_t2_croprs/checkpoints/epoch=000070-step=000219999.ckpt' 
 
-    config_file = args.config_file
-    weight_file = args.weight_file
+    config_file = args.config
+    weight_file = args.weight
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     if weight_file.startswith('http'):  # download from url
         weight_dir = Path('./weights')

cldm/transformations2.py

@@ -367,7 +367,7 @@ class TransformNet(nn.Module):
     def apply_transform_on_pil_image(self, x, tform_name):
         # x: PIL image
         # return: PIL image
-        assert tform_name in self.optional_names + ['Random Crop', 'Random Flip']
+        assert tform_name in self.optional_names + ['Fixed Augment']
         # if tform_name == 'Random Crop':  # the only transform dependent on image size
         #     # crop equivalent to 224/256
         #     w, h = x.size
@@ -392,10 +392,11 @@ class TransformNet(nn.Module):
         x = x.resize((256, 256), Image.BILINEAR)
         x = np.array(x).astype(np.float32) / 255.  # [0, 255] -> [0, 1]
         x = torch.from_numpy(x).permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]
-        if tform_name == 'Random Flip':
-            x = self.fixed_transforms[0](x)
-        elif tform_name == 'Random Crop':
-            x = self.fixed_transforms[1](x)
+        if tform_name == 'Fixed Augment':
+            for tform in self.fixed_transforms:
+                x = tform(x)
+                if isinstance(x, tuple):
+                    x = x[0]
         else:
             tform_id = self.optional_names.index(tform_name)
             tform = self.optional_transforms[tform_id]

pages/Extract_Secret.py

@@ -27,17 +27,17 @@ from io import BytesIO
 from tools.helpers import welcome_message
 from tools.ecc import BCH, RSC
 import streamlit as st
-from Embed_Secret import load_ecc, load_model, decode_secret, to_bytes, model_names, SECRET_LEN
+from Embed_Secret import parse_st_args, load_ecc, load_model, decode_secret, to_bytes, model_names, SECRET_LEN
 
 
 # model_names = ['RoSteALS', 'UNet']
 # SECRET_LEN = 100
 
-def app():
+def app(args):
     st.title('Watermarking Demo')
     # setup model
     model_name = st.selectbox("Choose the model", model_names)
-    model, tform_emb, tform_det = load_model(model_name)
+    model, tform_emb, tform_det = load_model(model_name, args)
     display_width = 300
     ecc = load_ecc('BCH')
     noise = TransformNet(p=1.0, crop_mode='resized_crop')
@@ -55,11 +55,11 @@ def app():
     
     # add crop
     st.subheader("Corruptions")
-    crop_button = st.button('Regenerate Crop', key='crop')
+    crop_button = st.button('Regenerate Crop/Flip/Resize', key='crop')
     if image_file is not None:
-        im_crop = noise.apply_transform_on_pil_image(im, 'Random Crop')
+        im_crop = noise.apply_transform_on_pil_image(im, 'Fixed Augment')
         if crop_button:
-            im_crop = noise.apply_transform_on_pil_image(im, 'Random Crop')
+            im_crop = noise.apply_transform_on_pil_image(im, 'Fixed Augment')
         # st.image(im_crop, width=display_width)
     
     # add noise source 1
@@ -106,5 +106,6 @@ def app():
         bit_acc_status.markdown(f'Bit Accuracy: **{bit_acc*100:.2f}%**<br />Word Accuracy: **{word_acc}**', unsafe_allow_html=True)
 
 if __name__ == '__main__':
-    app()
+    args = parse_st_args()
+    app(args)