add 160bit support

Embed_Secret.py

@@ -32,7 +32,6 @@ from streamlit.source_util import (
 )
 
 model_names = ['UNet']
-SECRET_LEN = 100
 
 
 def delete_page(main_script_path_str, page_name):
@@ -110,8 +109,6 @@ def load_UNet(args):
 
     config = OmegaConf.load(config_file).model
     secret_len = config.params.secret_len
-    global SECRET_LEN
-    SECRET_LEN = secret_len
     print(f'Secret length: {secret_len}')
     model = instantiate_from_config(config)
     state_dict = torch.load(weight_file, map_location=torch.device('cpu'))
@@ -124,7 +121,7 @@ def load_UNet(args):
     print(f'Missed keys: {misses}\nIgnore keys: {ignores}')
     model = model.to(device)
     model.eval()
-    return model
+    return model, secret_len
 
 def embed_secret(model_name, model, cover, tform, secret):
     if model_name == 'UNet':
@@ -167,17 +164,19 @@ def load_model(model_name, _args):
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
         ])
-        model = load_UNet(_args)
+        model, secret_len = load_UNet(_args)
     else:
         raise NotImplementedError
-    return model, tform_emb, tform_det
+    return model, tform_emb, tform_det, secret_len
 
 
 @st.cache_resource
-def load_ecc(ecc_name):
+def load_ecc(ecc_name, secret_len):
     if ecc_name == 'BCH':
-        # ecc = BCH(285, 10, SECRET_LEN, verbose=True)
-        ecc = BCH(payload_len= SECRET_LEN, verbose=True)
+        if secret_len == 160:
+            ecc = BCH(285, 10, secret_len, verbose=True)
+        elif secret_len == 100:
+            ecc = BCH(137, 5, payload_len= secret_len, verbose=True)
     elif ecc_name == 'RSC':
         ecc = RSC(data_bytes=16, ecc_bytes=4, verbose=True)
     return ecc
@@ -213,12 +212,10 @@ 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, args)
+    model, tform_emb, tform_det, secret_len = load_model(model_name, args)
     display_width = 300
-
     # ecc
-    ecc = load_ecc('BCH')
-    assert ecc.get_total_len() == SECRET_LEN
+    ecc = load_ecc('BCH', secret_len)
 
     # setup st
     st.subheader("Input")

pages/Extract_Secret.py

@@ -27,19 +27,16 @@ from io import BytesIO
 from tools.helpers import welcome_message
 from tools.ecc import BCH, RSC
 import streamlit as st
-from Embed_Secret import parse_st_args, 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
 
 
-# model_names = ['RoSteALS', 'UNet']
-# SECRET_LEN = 100
-
 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, args)
+    model, tform_emb, tform_det, secret_len = load_model(model_name, args)
     display_width = 300
-    ecc = load_ecc('BCH')
+    ecc = load_ecc('BCH', secret_len)
     noise = TransformNet(p=1.0, crop_mode='resized_crop')
     noise_names = noise.optional_names