Upload folder using huggingface_hub

decoder.py

@@ -5,8 +5,6 @@ This model contains BOTH encoder and decoder in one checkpoint
 import tensorflow as tf
 import numpy as np
 from PIL import Image as PILImage
-from tensorflow.python.saved_model import tag_constants
-from tensorflow.python.saved_model import signature_constants
 
 
 class TFStegaStampDecoderPretrained:
@@ -24,31 +22,21 @@ class TFStegaStampDecoderPretrained:
         """
         print(f"Loading TensorFlow pretrained model from: {model_path}")
         
-        # Create session and load model
-        self.sess = tf.compat.v1.Session(graph=tf.Graph())
+        # Load model with TF2 API
+        loaded = tf.saved_model.load(model_path)
         
-        # Load the model
-        model = tf.compat.v1.saved_model.loader.load(
-            self.sess,
-            [tag_constants.SERVING],
-            model_path
-        )
+        # Get the serving function
+        self._infer_fn = loaded.signatures['serving_default']
         
-        # Get signature
-        signature_def = model.signature_def[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+        # Get input/output info
+        input_signature = self._infer_fn.structured_input_signature[1]
+        output_signature = self._infer_fn.structured_outputs
         
-        # Get input tensor (image)
-        input_image_name = signature_def.inputs['image'].name
-        self.input_image = self.sess.graph.get_tensor_by_name(input_image_name)
-        print(f"  Input tensor: {input_image_name}")
-        print(f"    Shape: {self.input_image.shape}")
-        
-        # Get output tensor (decoded secret)
-        output_secret_name = signature_def.outputs['decoded'].name
-        self.output_secret = self.sess.graph.get_tensor_by_name(output_secret_name)
-        print(f"  Output tensor: {output_secret_name}")
-        print(f"    Shape: {self.output_secret.shape}")
+        self._input_names = list(input_signature.keys())
+        self._output_names = list(output_signature.keys())
         
+        print(f"  Input tensors: {self._input_names}")
+        print(f"  Output tensors: {self._output_names}")
         print("TensorFlow decoder loaded successfully!")
     
     def decode(self, image):
@@ -84,12 +72,29 @@ class TFStegaStampDecoderPretrained:
             image = np.array(image_pil).astype(np.float32) / 255.0
             image = np.expand_dims(image, axis=0)
         
-        # Run inference
-        feed_dict = {self.input_image: image}
-        secret = self.sess.run(self.output_secret, feed_dict=feed_dict)
+        # Ensure float32
+        image = image.astype(np.float32)
+        
+        # Run inference with TF2 API
+        # Convert to tensor
+        image_tensor = tf.convert_to_tensor(image)
+        
+        # Call the function - use the correct input name from signature
+        if 'image' in self._input_names:
+            result = self._infer_fn(image=image_tensor)
+        else:
+            # Use the first input name
+            result = self._infer_fn(**{self._input_names[0]: image_tensor})
+        
+        # Get decoded output - use the correct output name
+        if 'decoded' in self._output_names:
+            decoded = result['decoded'].numpy()
+        else:
+            # Use the first output
+            decoded = list(result.values())[0].numpy()
         
         # Output is (batch, 100)
-        secret = secret[0]  # Remove batch dimension
+        secret = decoded[0]  # Remove batch dimension
         
         # Threshold to get binary values
         bits = (secret > 0.5).astype(np.float32)
@@ -101,9 +106,9 @@ class TFStegaStampDecoderPretrained:
         return self.decode(image)
     
     def close(self):
-        """Close the TensorFlow session"""
-        if hasattr(self, 'sess') and self.sess is not None:
-            self.sess.close()
+        """Cleanup resources"""
+        # TF2 models don't need explicit session closing
+        pass
     
     def __del__(self):
         """Cleanup on deletion"""