quantumcontrol/stable-video-diffusion

1import argparse2 3import cv24import numpy as np5 6try:7    from imwatermark import WatermarkDecoder8except ImportError as e:9    try:10        # Assume some of the other dependencies such as torch are not fulfilled11        # import file without loading unnecessary libraries.12        import importlib.util13        import sys14 15        spec = importlib.util.find_spec("imwatermark.maxDct")16        assert spec is not None17        maxDct = importlib.util.module_from_spec(spec)18        sys.modules["maxDct"] = maxDct19        spec.loader.exec_module(maxDct)20 21        class WatermarkDecoder(object):22            """A minimal version of23            https://github.com/ShieldMnt/invisible-watermark/blob/main/imwatermark/watermark.py24            to only reconstruct bits using dwtDct"""25 26            def __init__(self, wm_type="bytes", length=0):27                assert wm_type == "bits", "Only bits defined in minimal import"28                self._wmType = wm_type29                self._wmLen = length30 31            def reconstruct(self, bits):32                if len(bits) != self._wmLen:33                    raise RuntimeError("bits are not matched with watermark length")34 35                return bits36 37            def decode(self, cv2Image, method="dwtDct", **configs):38                (r, c, channels) = cv2Image.shape39                if r * c < 256 * 256:40                    raise RuntimeError("image too small, should be larger than 256x256")41 42                bits = []43                assert method == "dwtDct"44                embed = maxDct.EmbedMaxDct(watermarks=[], wmLen=self._wmLen, **configs)45                bits = embed.decode(cv2Image)46                return self.reconstruct(bits)47 48    except:49        raise e50 51 52# A fixed 48-bit message that was choosen at random53# WATERMARK_MESSAGE = 0xB3EC907BB19E54WATERMARK_MESSAGE = 0b10110011111011001001000001111011101100011001111055# bin(x)[2:] gives bits of x as str, use int to convert them to 0/156WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]57MATCH_VALUES = [58    [27, "No watermark detected"],59    [33, "Partial watermark match. Cannot determine with certainty."],60    [61        35,62        (63            "Likely watermarked. In our test 0.02% of real images were "64            'falsely detected as "Likely watermarked"'65        ),66    ],67    [68        49,69        (70            "Very likely watermarked. In our test no real images were "71            'falsely detected as "Very likely watermarked"'72        ),73    ],74]75 76 77class GetWatermarkMatch:78    def __init__(self, watermark):79        self.watermark = watermark80        self.num_bits = len(self.watermark)81        self.decoder = WatermarkDecoder("bits", self.num_bits)82 83    def __call__(self, x: np.ndarray) -> np.ndarray:84        """85        Detects the number of matching bits the predefined watermark with one86        or multiple images. Images should be in cv2 format, e.g. h x w x c BGR.87 88        Args:89            x: ([B], h w, c) in range [0, 255]90 91        Returns:92           number of matched bits ([B],)93        """94        squeeze = len(x.shape) == 395        if squeeze:96            x = x[None, ...]97 98        bs = x.shape[0]99        detected = np.empty((bs, self.num_bits), dtype=bool)100        for k in range(bs):101            detected[k] = self.decoder.decode(x[k], "dwtDct")102        result = np.sum(detected == self.watermark, axis=-1)103        if squeeze:104            return result[0]105        else:106            return result107 108 109get_watermark_match = GetWatermarkMatch(WATERMARK_BITS)110 111 112if __name__ == "__main__":113    parser = argparse.ArgumentParser()114    parser.add_argument(115        "filename",116        nargs="+",117        type=str,118        help="Image files to check for watermarks",119    )120    opts = parser.parse_args()121 122    print(123        """124        This script tries to detect watermarked images. Please be aware of125        the following:126        - As the watermark is supposed to be invisible, there is the risk that127          watermarked images may not be detected.128        - To maximize the chance of detection make sure that the image has the same129          dimensions as when the watermark was applied (most likely 1024x1024130          or 512x512).131        - Specific image manipulation may drastically decrease the chance that132          watermarks can be detected.133        - There is also the chance that an image has the characteristics of the134          watermark by chance.135        - The watermark script is public, anybody may watermark any images, and136          could therefore claim it to be generated.137        - All numbers below are based on a test using 10,000 images without any138          modifications after applying the watermark.139        """140    )141 142    for fn in opts.filename:143        image = cv2.imread(fn)144        if image is None:145            print(f"Couldn't read {fn}. Skipping")146            continue147 148        num_bits = get_watermark_match(image)149        k = 0150        while num_bits > MATCH_VALUES[k][0]:151            k += 1152        print(153            f"{fn}: {MATCH_VALUES[k][1]}",154            f"Bits that matched the watermark {num_bits} from {len(WATERMARK_BITS)}\n",155            sep="\n\t",156        )157