Skip to content

End2EndStegNet

stegobox.codec.End2EndStegNet

Bases: BaseCodec

End2EndStegNet is a steganography method for Image.

Source code derived from: qzramiz/End-To-End-Image-Steganography

Source code in stegobox/codec/end2end_stegnet/end2end_stegnet.py
class End2EndStegNet(BaseCodec):
    """
    End2EndStegNet is a steganography method for Image.

    Source code derived from:
    [qzramiz/End-To-End-Image-Steganography](https://github.com/qzramiz/End-To-End-Image-Steganography)
    """

    def __init__(self, weight_file="ckpt/end2end/end2end.pickle") -> None:
        super().__init__()
        self.secret = torch.zeros(1, 1)
        self.stego = torch.zeros(1, 1)
        # self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device("cpu")
        self.model = EDS().to(self.device)
        self.weight_file = weight_file

        weights = torch.load(self.weight_file, map_location=self.device)
        self.model.load_state_dict(weights)
        self.model.eval()

    def _transform(self, image: Image.Image) -> torch.Tensor:
        transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Resize((64, 64))]
        )
        return transform(image)  # type: ignore

    def encode(self, carrier: Image.Image, payload: Image.Image) -> Image.Image:
        """Merge secret_img into host image

        Args:
            carrier: host image
            payload: payload in bits per pixel

        Returns:
            Image: stego image.
        """

        cover = self._transform(carrier)
        cover = cover.unsqueeze(0)

        payload_t = rgb_to_grayscale(self._transform(payload))
        payload_t = payload_t.unsqueeze(0)
        # print(cover.shape, payload_t.shape)

        self.stego, self.secret = self.model(cover, payload_t)
        return to_pil_image(self.stego.squeeze(0))

    def decode(self, carrier: Image.Image) -> Image.Image:
        """Extract payload from stego image

        Args:
            carrier: stego image

        Returns:
            Decoded payload image
        """

        return to_pil_image(self.secret.squeeze(0))

encode(carrier, payload)

Merge secret_img into host image

Parameters:

Name Type Description Default
carrier Image

host image

required
payload Image

payload in bits per pixel

required

Returns:

Name Type Description
Image Image

stego image.

Source code in stegobox/codec/end2end_stegnet/end2end_stegnet.py
def encode(self, carrier: Image.Image, payload: Image.Image) -> Image.Image:
    """Merge secret_img into host image

    Args:
        carrier: host image
        payload: payload in bits per pixel

    Returns:
        Image: stego image.
    """

    cover = self._transform(carrier)
    cover = cover.unsqueeze(0)

    payload_t = rgb_to_grayscale(self._transform(payload))
    payload_t = payload_t.unsqueeze(0)
    # print(cover.shape, payload_t.shape)

    self.stego, self.secret = self.model(cover, payload_t)
    return to_pil_image(self.stego.squeeze(0))

decode(carrier)

Extract payload from stego image

Parameters:

Name Type Description Default
carrier Image

stego image

required

Returns:

Type Description
Image

Decoded payload image

Source code in stegobox/codec/end2end_stegnet/end2end_stegnet.py
def decode(self, carrier: Image.Image) -> Image.Image:
    """Extract payload from stego image

    Args:
        carrier: stego image

    Returns:
        Decoded payload image
    """

    return to_pil_image(self.secret.squeeze(0))