Skip to content

RivaGAN

stegobox.codec.RivaGAN

Bases: BaseCodec

This implementation of the RivaGAN is from ShieldMnt/invisible-watermark which provides a pretrained onnx model for the RivaGAN encoder and decoder.

The original implementation of RivaGAN is from source code: DAI-Lab/RivaGAN and paper: Robust Invisible Video Watermarking with Attention - https://arxiv.org/abs/1909.01285.

Tip: Encoded payload of this method is limited to 32 bytes.

Source code in stegobox/codec/invisible_watermark/riva_gan.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class RivaGAN(BaseCodec):
    """This implementation of the RivaGAN is from
    [ShieldMnt/invisible-watermark](https://github.com/ShieldMnt/invisible-watermark)
    which provides a pretrained onnx model for the RivaGAN encoder and decoder.

    The original implementation of RivaGAN is from source code:
    [DAI-Lab/RivaGAN](https://github.com/DAI-Lab/RivaGAN) and paper: *Robust Invisible
    Video Watermarking with Attention* - https://arxiv.org/abs/1909.01285.

    Tip: Encoded payload of this method is limited to 32 bytes.
    """

    def __init__(
        self,
        threshold: float = 0.52,
        encoder_onnx: str = "ckpt/riva_gan/rivagan_encoder.onnx",
        decoder_onnx: str = "ckpt/riva_gan/rivagan_decoder.onnx",
    ) -> None:
        super().__init__()

        self.threshold = threshold
        self.encoder = onnxruntime.InferenceSession(encoder_onnx)
        self.decoder = onnxruntime.InferenceSession(decoder_onnx)

    def encode(self, carrier: cv2.Mat, payload: str) -> cv2.Mat:
        """Encodes payload string into a frame of a video with a RivaGAN encoder.

        Args:
            carrier: The carrier frame, read with `io.image.read_cv2()`.
            payload: The payload string.

        Returns:
            The encoded image, write with `io.image.write_cv2()`.
        """
        frame = torch.from_numpy(np.array([carrier], dtype=np.float32)) / 127.5 - 1.0
        frame = frame.permute(3, 0, 1, 2).unsqueeze(0)

        payload_seq = payload_in_bytes(payload)
        # Only first 32 bytes supported for RivaGAN encoding
        payload_seq = payload_seq[:32]

        data = torch.from_numpy(np.array([payload_seq], dtype=np.float32))

        inputs = {
            "frame": frame.detach().cpu().numpy(),
            "data": data.detach().cpu().numpy(),
        }
        outputs = self.encoder.run(None, inputs)

        wm_frame = outputs[0]
        wm_frame = torch.clamp(torch.from_numpy(wm_frame), min=-1.0, max=1.0)
        wm_frame = (
            ((wm_frame[0, :, 0, :, :].permute(1, 2, 0) + 1.0) * 127.5)
            .detach()
            .cpu()
            .numpy()
            .astype("uint8")
        )

        return wm_frame

    def decode(self, carrier: cv2.Mat) -> str:
        """Decodes payload string from a frame of a video with a RivaGAN decoder.

        Args:
            carrier: The carrier frame, read with `io.image.read_cv2()`.

        Returns:
            The decoded payload string.
        """
        frame = torch.from_numpy(np.array([carrier], dtype=np.float32)) / 127.5 - 1.0
        frame = frame.permute(3, 0, 1, 2).unsqueeze(0)

        inputs = {
            "frame": frame.detach().cpu().numpy(),
        }
        outputs = self.decoder.run(None, inputs)

        data = outputs[0][0]
        decoded = np.array(data > self.threshold, dtype=np.uint8)
        return reconstruct_bytes(decoded, 32)

encode(carrier, payload)

Encodes payload string into a frame of a video with a RivaGAN encoder.

Parameters:

Name Type Description Default
carrier Mat

The carrier frame, read with io.image.read_cv2().

 required
payload str

The payload string.

 required

Returns:

Type Description
Mat

The encoded image, write with io.image.write_cv2().
Source code in stegobox/codec/invisible_watermark/riva_gan.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def encode(self, carrier: cv2.Mat, payload: str) -> cv2.Mat:
    """Encodes payload string into a frame of a video with a RivaGAN encoder.

    Args:
        carrier: The carrier frame, read with `io.image.read_cv2()`.
        payload: The payload string.

    Returns:
        The encoded image, write with `io.image.write_cv2()`.
    """
    frame = torch.from_numpy(np.array([carrier], dtype=np.float32)) / 127.5 - 1.0
    frame = frame.permute(3, 0, 1, 2).unsqueeze(0)

    payload_seq = payload_in_bytes(payload)
    # Only first 32 bytes supported for RivaGAN encoding
    payload_seq = payload_seq[:32]

    data = torch.from_numpy(np.array([payload_seq], dtype=np.float32))

    inputs = {
        "frame": frame.detach().cpu().numpy(),
        "data": data.detach().cpu().numpy(),
    }
    outputs = self.encoder.run(None, inputs)

    wm_frame = outputs[0]
    wm_frame = torch.clamp(torch.from_numpy(wm_frame), min=-1.0, max=1.0)
    wm_frame = (
        ((wm_frame[0, :, 0, :, :].permute(1, 2, 0) + 1.0) * 127.5)
        .detach()
        .cpu()
        .numpy()
        .astype("uint8")
    )

    return wm_frame

decode(carrier)

Decodes payload string from a frame of a video with a RivaGAN decoder.

Parameters:

Name Type Description Default
carrier Mat

The carrier frame, read with io.image.read_cv2().

 required

Returns:

Type Description
str

The decoded payload string.
Source code in stegobox/codec/invisible_watermark/riva_gan.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def decode(self, carrier: cv2.Mat) -> str:
    """Decodes payload string from a frame of a video with a RivaGAN decoder.

    Args:
        carrier: The carrier frame, read with `io.image.read_cv2()`.

    Returns:
        The decoded payload string.
    """
    frame = torch.from_numpy(np.array([carrier], dtype=np.float32)) / 127.5 - 1.0
    frame = frame.permute(3, 0, 1, 2).unsqueeze(0)

    inputs = {
        "frame": frame.detach().cpu().numpy(),
    }
    outputs = self.decoder.run(None, inputs)

    data = outputs[0][0]
    decoded = np.array(data > self.threshold, dtype=np.uint8)
    return reconstruct_bytes(decoded, 32)