DAI-Lab / SteganoGAN

SteganoGAN is a tool for creating steganographic images using adversarial training.
MIT License
305 stars 85 forks source link

How can I train this model on different images ? #69

Open purushottam22 opened 3 years ago

purushottam22 commented 3 years ago

Description

I tried to run this model on different images (like leena). I get segmentation fault.

I want to train this model on different images.

kompowiec commented 2 weeks ago

you have in tutorial... example code (make sure add "default" folder inside).

# Imports
import numpy as np
from steganogan import SteganoGAN
from steganogan.loader import DataLoader
from steganogan.encoders import BasicEncoder, DenseEncoder
from steganogan.decoders import BasicDecoder, DenseDecoder
from steganogan.critics import BasicCritic

# Load Data
train = DataLoader('/home/$USER/wallpapers', limit=np.inf, shuffle=True, batch_size=3)
validation = DataLoader('/home/$USER/wallpapers', limit=np.inf, shuffle=True, batch_size=3)

# Create SteganoGAN Instance
steganogan = SteganoGAN(
    data_depth=1,                   # Number of layers for representing data
    encoder=BasicEncoder,           # Choose between BasicEncoder or DenseEncoder
    decoder=BasicDecoder,           # Choose between BasicDecoder or DenseDecoder
    critic=BasicCritic,             # BasicCritic is used here
    hidden_size=32,                 # Number of channels for hidden layers
    cuda=True,                      # Enable CUDA if available
    verbose=True                    # Print training info to console
)

# Train and Save the Model
steganogan.fit(train, validation, epochs=99)
steganogan.save('demo.steg')  # Save model as 'demo.steg'

# Load the Model
# If you have a saved model and want to load it:
steganogan = SteganoGAN.load(architecture='basic', path='demo.steg', cuda=True, verbose=True)

# Encoding
input_image = 'input.png'
output_image = 'output.png'
secret_message = 'This is a super secret message!'
steganogan.encode(input_image, output_image, secret_message)
print("Encoding completed.")

# Decoding
decoded_message = steganogan.decode(output_image)
print("Decoded message:", decoded_message)