pseudopros 0.1.0__py3-none-any.whl

pseudopros/__init__.py

@@ -0,0 +1,26 @@
+from pseudopros.anonymizer import Anonymizer
+from pseudopros.keying import PseudonymKey
+from pseudopros.security import SecurityAudit
+from pseudopros.synthesizer import (
+    NoiseSynthesizer,
+    RapidSynthesizer,
+    RosterSynthesizer,
+    build_face_bank,
+)
+
+try:
+    from importlib.metadata import metadata as _metadata, version as _version
+    __version__ = _version('pseudopros')
+    __doc__ = _metadata('pseudopros')['Description']
+except Exception:
+    __version__ = "0.0.0"
+
+__all__ = [
+    "Anonymizer",
+    "NoiseSynthesizer",
+    "RosterSynthesizer",
+    "PseudonymKey",
+    "RapidSynthesizer",
+    "SecurityAudit",
+    "build_face_bank",
+]

pseudopros/_cli.py

@@ -0,0 +1,159 @@
+"""Command-line interface for pseudopros."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import cv2
+import imageio
+
+from pseudopros.anonymizer import Anonymizer
+from pseudopros.synthesizer import NoiseSynthesizer, RosterSynthesizer
+
+
+def _load_secret(args: argparse.Namespace) -> bytes:
+    if args.secret_file:
+        return Path(args.secret_file).read_bytes()
+    if args.secret_hex:
+        return bytes.fromhex(args.secret_hex)
+    env = os.environ.get("PSEUDOPROS_SECRET")
+    if env:
+        return bytes.fromhex(env)
+    print(
+        "error: provide a secret via --secret-file, --secret-hex, or "
+        "PSEUDOPROS_SECRET env var",
+        file=sys.stderr,
+    )
+    sys.exit(1)
+
+
+def _load_salt(args: argparse.Namespace) -> bytes:
+    if args.salt_hex:
+        return bytes.fromhex(args.salt_hex)
+    env = os.environ.get("PSEUDOPROS_SALT")
+    if env:
+        return bytes.fromhex(env)
+    return b"pseudopros"
+
+
+def anonymize_video(
+    input_path: str,
+    output_path: str,
+    anon,
+    enrollment_path: str,
+) -> int:
+    """Process a video file frame-by-frame. Returns frame count processed."""
+    enroll_frame = imageio.v2.imread(enrollment_path)
+    if not anon.enroll(enroll_frame, "subject"):
+        print(
+            f"warning: no face detected in enrollment image {enrollment_path!r}",
+            file=sys.stderr,
+        )
+
+    cap = cv2.VideoCapture(input_path)
+    if not cap.isOpened():
+        print(f"error: cannot open input {input_path!r}", file=sys.stderr)
+        sys.exit(1)
+
+    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    count = 0
+    while True:
+        ret, frame_bgr = cap.read()
+        if not ret:
+            break
+        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+        result_rgb = anon.process(frame_rgb)
+        result_bgr = cv2.cvtColor(result_rgb, cv2.COLOR_RGB2BGR)
+        writer.write(result_bgr)
+        count += 1
+
+    cap.release()
+    writer.release()
+    return count
+
+
+def _cmd_anonymize(args: argparse.Namespace) -> None:
+    secret = _load_secret(args)
+    salt = _load_salt(args)
+    synth = RosterSynthesizer() if args.synthesizer == "roster" else NoiseSynthesizer()
+
+    print(
+        f"Loading FOMM model {'from ' + args.checkpoint if args.checkpoint else '(downloading from HuggingFace)'}...",
+        file=sys.stderr,
+    )
+    anon = Anonymizer.build(
+        secret=secret,
+        checkpoint_path=args.checkpoint or None,
+        config_path=args.config or None,
+        salt=salt,
+        device=args.device,
+        synthesizer=synth,
+    )
+
+    print(f"Processing {args.input!r} -> {args.output!r} ...", file=sys.stderr)
+    count = anonymize_video(args.input, args.output, anon, args.enrollment)
+    print(f"Done. {count} frame(s) written to {args.output!r}", file=sys.stderr)
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="pseudopros",
+        description="Cryptographic face anonymization for video pipelines.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    anon_p = sub.add_parser("anonymize", help="Anonymize faces in a video file.")
+    anon_p.add_argument("input", help="Input video path.")
+    anon_p.add_argument("output", help="Output video path.")
+    anon_p.add_argument(
+        "--enrollment", required=True, metavar="IMAGE",
+        help="Image file containing the subject's face for enrollment.",
+    )
+
+    secret_g = anon_p.add_mutually_exclusive_group()
+    secret_g.add_argument(
+        "--secret-file", metavar="PATH",
+        help="Path to a file containing the raw secret bytes. Recommended for production.",
+    )
+    secret_g.add_argument(
+        "--secret-hex", metavar="HEX",
+        help="Secret key as a hex string. Or set PSEUDOPROS_SECRET env var.",
+    )
+
+    anon_p.add_argument(
+        "--salt-hex", metavar="HEX",
+        help="Deployment salt as a hex string. Or set PSEUDOPROS_SALT env var. "
+             "Default: hex of b'pseudopros'.",
+    )
+    anon_p.add_argument(
+        "--synthesizer", choices=["noise", "roster"], default="noise",
+        help="Synthesizer backend. 'noise' is a placeholder; "
+             "'roster' uses LDM (~1.3 GB download on first use). Default: noise.",
+    )
+    anon_p.add_argument("--device", default="cpu", help="Torch device. Default: cpu.")
+    anon_p.add_argument(
+        "--checkpoint", metavar="PATH",
+        help="Path to a local FOMM checkpoint (.pth.tar). "
+             "If omitted, weights are downloaded from HuggingFace Hub.",
+    )
+    anon_p.add_argument(
+        "--config", metavar="PATH",
+        help="Path to the FOMM config YAML. Required when --checkpoint is set.",
+    )
+
+    args = parser.parse_args(argv)
+    if args.command == "anonymize":
+        _cmd_anonymize(args)
+
+
+if __name__ == "__main__":
+    main()

pseudopros/_fomm/modules/dense_motion.py

@@ -0,0 +1,113 @@
+from torch import nn
+import torch.nn.functional as F
+import torch
+from .util import Hourglass, AntiAliasInterpolation2d, make_coordinate_grid, kp2gaussian
+
+
+class DenseMotionNetwork(nn.Module):
+    """
+    Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving
+    """
+
+    def __init__(self, block_expansion, num_blocks, max_features, num_kp, num_channels, estimate_occlusion_map=False,
+                 scale_factor=1, kp_variance=0.01):
+        super(DenseMotionNetwork, self).__init__()
+        self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp + 1) * (num_channels + 1),
+                                   max_features=max_features, num_blocks=num_blocks)
+
+        self.mask = nn.Conv2d(self.hourglass.out_filters, num_kp + 1, kernel_size=(7, 7), padding=(3, 3))
+
+        if estimate_occlusion_map:
+            self.occlusion = nn.Conv2d(self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3))
+        else:
+            self.occlusion = None
+
+        self.num_kp = num_kp
+        self.scale_factor = scale_factor
+        self.kp_variance = kp_variance
+
+        if self.scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
+
+    def create_heatmap_representations(self, source_image, kp_driving, kp_source):
+        """
+        Eq 6. in the paper H_k(z)
+        """
+        spatial_size = source_image.shape[2:]
+        gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance)
+        gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance)
+        heatmap = gaussian_driving - gaussian_source
+
+        #adding background feature
+        zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]).type(heatmap.type())
+        heatmap = torch.cat([zeros, heatmap], dim=1)
+        heatmap = heatmap.unsqueeze(2)
+        return heatmap
+
+    def create_sparse_motions(self, source_image, kp_driving, kp_source):
+        """
+        Eq 4. in the paper T_{s<-d}(z)
+        """
+        bs, _, h, w = source_image.shape
+        identity_grid = make_coordinate_grid((h, w), type=kp_source['value'].type())
+        identity_grid = identity_grid.view(1, 1, h, w, 2)
+        coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2)
+        if 'jacobian' in kp_driving:
+            jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian']))
+            jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
+            jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
+            coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
+            coordinate_grid = coordinate_grid.squeeze(-1)
+
+        driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2)
+
+        #adding background feature
+        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
+        sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1)
+        return sparse_motions
+
+    def create_deformed_source_image(self, source_image, sparse_motions):
+        r"""
+        Eq 7. in the paper \hat{T}_{s<-d}(z)
+        """
+        bs, _, h, w = source_image.shape
+        source_repeat = source_image.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp + 1, 1, 1, 1, 1)
+        source_repeat = source_repeat.view(bs * (self.num_kp + 1), -1, h, w)
+        sparse_motions = sparse_motions.view((bs * (self.num_kp + 1), h, w, -1))
+        sparse_deformed = F.grid_sample(source_repeat, sparse_motions, align_corners=True)
+        sparse_deformed = sparse_deformed.view((bs, self.num_kp + 1, -1, h, w))
+        return sparse_deformed
+
+    def forward(self, source_image, kp_driving, kp_source):
+        if self.scale_factor != 1:
+            source_image = self.down(source_image)
+
+        bs, _, h, w = source_image.shape
+
+        out_dict = dict()
+        heatmap_representation = self.create_heatmap_representations(source_image, kp_driving, kp_source)
+        sparse_motion = self.create_sparse_motions(source_image, kp_driving, kp_source)
+        deformed_source = self.create_deformed_source_image(source_image, sparse_motion)
+        out_dict['sparse_deformed'] = deformed_source
+
+        input = torch.cat([heatmap_representation, deformed_source], dim=2)
+        input = input.view(bs, -1, h, w)
+
+        prediction = self.hourglass(input)
+
+        mask = self.mask(prediction)
+        mask = F.softmax(mask, dim=1)
+        out_dict['mask'] = mask
+        mask = mask.unsqueeze(2)
+        sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3)
+        deformation = (sparse_motion * mask).sum(dim=1)
+        deformation = deformation.permute(0, 2, 3, 1)
+
+        out_dict['deformation'] = deformation
+
+        # Sec. 3.2 in the paper
+        if self.occlusion:
+            occlusion_map = torch.sigmoid(self.occlusion(prediction))
+            out_dict['occlusion_map'] = occlusion_map
+
+        return out_dict

pseudopros/_fomm/modules/generator.py

@@ -0,0 +1,97 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d
+from .dense_motion import DenseMotionNetwork
+
+
+class OcclusionAwareGenerator(nn.Module):
+    """
+    Generator that given source image and and keypoints try to transform image according to movement trajectories
+    induced by keypoints. Generator follows Johnson architecture.
+    """
+
+    def __init__(self, num_channels, num_kp, block_expansion, max_features, num_down_blocks,
+                 num_bottleneck_blocks, estimate_occlusion_map=False, dense_motion_params=None, estimate_jacobian=False):
+        super(OcclusionAwareGenerator, self).__init__()
+
+        if dense_motion_params is not None:
+            self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, num_channels=num_channels,
+                                                           estimate_occlusion_map=estimate_occlusion_map,
+                                                           **dense_motion_params)
+        else:
+            self.dense_motion_network = None
+
+        self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3))
+
+        down_blocks = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2 ** i))
+            out_features = min(max_features, block_expansion * (2 ** (i + 1)))
+            down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        up_blocks = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i)))
+            out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1)))
+            up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
+        self.up_blocks = nn.ModuleList(up_blocks)
+
+        self.bottleneck = torch.nn.Sequential()
+        in_features = min(max_features, block_expansion * (2 ** num_down_blocks))
+        for i in range(num_bottleneck_blocks):
+            self.bottleneck.add_module('r' + str(i), ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1)))
+
+        self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3))
+        self.estimate_occlusion_map = estimate_occlusion_map
+        self.num_channels = num_channels
+
+    def deform_input(self, inp, deformation):
+        _, h_old, w_old, _ = deformation.shape
+        _, _, h, w = inp.shape
+        if h_old != h or w_old != w:
+            deformation = deformation.permute(0, 3, 1, 2)
+            deformation = F.interpolate(deformation, size=(h, w), mode='bilinear')
+            deformation = deformation.permute(0, 2, 3, 1)
+        return F.grid_sample(inp, deformation, align_corners=True)
+
+    def forward(self, source_image, kp_driving, kp_source):
+        # Encoding (downsampling) part
+        out = self.first(source_image)
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out)
+
+        # Transforming feature representation according to deformation and occlusion
+        output_dict = {}
+        if self.dense_motion_network is not None:
+            dense_motion = self.dense_motion_network(source_image=source_image, kp_driving=kp_driving,
+                                                     kp_source=kp_source)
+            output_dict['mask'] = dense_motion['mask']
+            output_dict['sparse_deformed'] = dense_motion['sparse_deformed']
+
+            if 'occlusion_map' in dense_motion:
+                occlusion_map = dense_motion['occlusion_map']
+                output_dict['occlusion_map'] = occlusion_map
+            else:
+                occlusion_map = None
+            deformation = dense_motion['deformation']
+            out = self.deform_input(out, deformation)
+
+            if occlusion_map is not None:
+                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
+                    occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
+                out = out * occlusion_map
+
+            output_dict["deformed"] = self.deform_input(source_image, deformation)
+
+        # Decoding part
+        out = self.bottleneck(out)
+        for i in range(len(self.up_blocks)):
+            out = self.up_blocks[i](out)
+        out = self.final(out)
+        out = F.sigmoid(out)
+
+        output_dict["prediction"] = out
+
+        return output_dict

pseudopros/_fomm/modules/keypoint_detector.py

@@ -0,0 +1,75 @@
+from torch import nn
+import torch
+import torch.nn.functional as F
+from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d
+
+
+class KPDetector(nn.Module):
+    """
+    Detecting a keypoints. Return keypoint position and jacobian near each keypoint.
+    """
+
+    def __init__(self, block_expansion, num_kp, num_channels, max_features,
+                 num_blocks, temperature, estimate_jacobian=False, scale_factor=1,
+                 single_jacobian_map=False, pad=0):
+        super(KPDetector, self).__init__()
+
+        self.predictor = Hourglass(block_expansion, in_features=num_channels,
+                                   max_features=max_features, num_blocks=num_blocks)
+
+        self.kp = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_kp, kernel_size=(7, 7),
+                            padding=pad)
+
+        if estimate_jacobian:
+            self.num_jacobian_maps = 1 if single_jacobian_map else num_kp
+            self.jacobian = nn.Conv2d(in_channels=self.predictor.out_filters,
+                                      out_channels=4 * self.num_jacobian_maps, kernel_size=(7, 7), padding=pad)
+            self.jacobian.weight.data.zero_()
+            self.jacobian.bias.data.copy_(torch.tensor([1, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float))
+        else:
+            self.jacobian = None
+
+        self.temperature = temperature
+        self.scale_factor = scale_factor
+        if self.scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
+
+    def gaussian2kp(self, heatmap):
+        """
+        Extract the mean and from a heatmap
+        """
+        shape = heatmap.shape
+        heatmap = heatmap.unsqueeze(-1)
+        grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
+        value = (heatmap * grid).sum(dim=(2, 3))
+        kp = {'value': value}
+
+        return kp
+
+    def forward(self, x):
+        if self.scale_factor != 1:
+            x = self.down(x)
+
+        feature_map = self.predictor(x)
+        prediction = self.kp(feature_map)
+
+        final_shape = prediction.shape
+        heatmap = prediction.view(final_shape[0], final_shape[1], -1)
+        heatmap = F.softmax(heatmap / self.temperature, dim=2)
+        heatmap = heatmap.view(*final_shape)
+
+        out = self.gaussian2kp(heatmap)
+
+        if self.jacobian is not None:
+            jacobian_map = self.jacobian(feature_map)
+            jacobian_map = jacobian_map.reshape(final_shape[0], self.num_jacobian_maps, 4, final_shape[2],
+                                                final_shape[3])
+            heatmap = heatmap.unsqueeze(2)
+
+            jacobian = heatmap * jacobian_map
+            jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1)
+            jacobian = jacobian.sum(dim=-1)
+            jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2)
+            out['jacobian'] = jacobian
+
+        return out