Re-implementation for ICCV23 "Social Diffusion: Long-term Multiple Human Motion Anticipation"
import os
from os.path import join, isdir
from social_diffusion.datasets.haggling import get_haggling_test_sequences
from social_diffusion.evaluation.db import Database
import numpy as np
from tqdm import tqdm
from social_diffusion import get_output_dir
from social_diffusion.eval_utils import load_trained_model, predict


def evaluation():
    INPUT_FRAME_10PERC = 178
    AVG_TEST_FRAMES = 1784

    ds_train, ema_diffusion = load_trained_model()

    test_seq = get_haggling_test_sequences(to_sequence=False)
    db = Database(test_seq, motion_word_size=10)

    eval_out_dir = join(get_output_dir(), "eval")
    if not isdir(eval_out_dir):
        os.makedirs(eval_out_dir)

    avg_ndms_per_seq = []
    for seqid, Seq in enumerate(get_haggling_test_sequences()):
        pred = predict(ema_diffusion=ema_diffusion, ds_train=ds_train, Seq=Seq)
        pred = pred[:, INPUT_FRAME_10PERC - 10 : AVG_TEST_FRAMES]  # noqa E203
        # pred --> n_samples x n_frames x n_person x Jd
        n_person = pred.shape[2]
        fname_out = join(eval_out_dir, f"ndms_seq{seqid}.npy")
        n_batch = pred.shape[0]
        Values = []
        for b in tqdm(range(n_batch), leave=True):
            for pid in range(n_person):
                seq = pred[b, :, pid]
                d, _ = db.rolling_query(seq)
                Values.append(d)
        Values = np.array(Values, dtype=np.float32)
        np.save(fname_out, Values)  # save per-frame ndms to file..

        avg_ndms_per_seq.append(np.mean(Values))
    print("Avg ndms: ", avg_ndms_per_seq)


if __name__ == "__main__":
    evaluation()

eval_haggling_ndms.py
jutanke/social_diffusion
from social_diffusion.datasets.haggling import (
    get_haggling_test_sequences,
    get_haggling_train_sequences,
)
import numpy as np
from einops import reduce
from tqdm import tqdm
import torch
from social_diffusion.evaluation.decider import WhoSpeaksDecider
from social_diffusion.evaluation.buyer_decider import WhoIsBuyerDecider
import social_diffusion.evaluation.jsd as jsd
from social_diffusion.eval_utils import load_trained_model, predict
from scipy.spatial.distance import jensenshannon
from typing import List

device = torch.device("cuda")
whospeaks = WhoSpeaksDecider(device=device)
whoisbuyer = WhoIsBuyerDecider(device=device)


def get_p_from_sequence(sequence: List[np.ndarray]):
    HM = []
    for Seq in tqdm(sequence):
        buyer_id = whoisbuyer.who_is_the_buyer(Seq)
        attn, seller_ids = jsd.calculate_attn(Seq, buyer_id=buyer_id)
        speech_all = whospeaks.eval(Seq)
        speech = speech_all[:, seller_ids]
        speech_buyer = speech_all[:, buyer_id]
        words = jsd.attn_to_word5(
            attn=attn, speech=speech, buyer_speech=speech_buyer
        )  # noqa E501
        hm = jsd.get_hm(words)
        HM.append(hm)

    HM = reduce(np.array(HM, dtype=np.float32), "c w h -> w h", "sum")
    return jsd.to_prob(HM)


def evaluation():
    INPUT_FRAME_10PERC = 178
    AVG_TEST_FRAMES = 1784

    ds_train, ema_diffusion = load_trained_model()

    TEST = [
        seq.Seq[INPUT_FRAME_10PERC:] for seq in get_haggling_test_sequences()
    ]  # noqa E501
    TRAIN = [
        seq.Seq[INPUT_FRAME_10PERC:] for seq in get_haggling_train_sequences()
    ]  # noqa E501

    p_test = get_p_from_sequence(TEST)
    p_train = get_p_from_sequence(TRAIN)

    print("train: JSD: ", jensenshannon(p_test, p_train))

    Pred = []
    for Seq in get_haggling_test_sequences():
        pred = predict(ema_diffusion=ema_diffusion, ds_train=ds_train, Seq=Seq)
        pred = pred[:, INPUT_FRAME_10PERC:AVG_TEST_FRAMES]  # noqa E203
        Pred.append(pred)
    Pred = np.concatenate(Pred, axis=0)  # (b s) t p Jd
    p_method = get_p_from_sequence(Pred)
    print("JSD: ", jensenshannon(p_test, p_method))


if __name__ == "__main__":
    evaluation()

eval_haggling_sscp.py
jutanke/social_diffusion
import torch
import os
from os.path import join, isdir
from os import listdir, makedirs
import shutil
from sys import platform


def create_training_dir(path: str):
    """
    :param path
    """
    if isdir(path):
        for elem in listdir(path):
            if not elem.endswith(".pt"):
                raise ValueError(
                    f"The path <{path}> has a file that seems suspicious for deletion: {elem}"  # noqa E501
                )
        shutil.rmtree(path)
    makedirs(path)


def get_number_of_processes():
    if platform == "linux" or platform == "linux2":
        # linux
        return 10
    elif platform == "darwin":
        # OS X
        return 0  # noqa E501 weird fork bug with <get_LocalMultiPersonDataset.<locals>.create_batch_entry_fn>
    elif platform == "win32":
        # Windows...
        return 10
    raise ValueError(f"cannot determine platform {platform}")


def get_device():
    if torch.backends.mps.is_available() and torch.backends.mps.is_built():
        return torch.device("mps")
    elif torch.cuda.is_available():
        return torch.device("cuda")
    else:
        return torch.device("cpu")


def get_cache_dir() -> str:
    cache_dir = join(os.getcwd(), ".cache")
    if not isdir(cache_dir):
        os.makedirs(cache_dir)
    return cache_dir


def get_data_dir() -> str:
    data_dir = join(os.getcwd(), "data")
    if not isdir(data_dir):
        raise ValueError(f"We require the data path to exist: {data_dir}")
    return data_dir


def get_video_dir(video_name) -> str:
    return join(get_output_dir(), f"vid/{video_name}")


def get_output_dir() -> str:
    cache_dir = join(os.getcwd(), "output")
    if not isdir(cache_dir):
        os.makedirs(cache_dir)
    return cache_dir


class WeirdShapeError(ValueError):
    def __init__(self, message):
        super().__init__(message)
social_diffusion/__init__.py
jutanke/social_diffusion
import json
from typing import List
import numpy as np
from einops import rearrange, reduce


def calculate_standardization(all_poses, skel):
    """
    :param all_poses: {b t p jd}
    """
    eps = 0.000000001
    all_poses = rearrange(all_poses, "b t p jd -> (b p t) jd")
    p_mu = reduce(all_poses, "t jd -> jd", "mean")
    p_std = np.std(all_poses, axis=0) + eps
    return p_mu, p_std


def load_json(fname):
    with open(fname, "r") as f:
        return json.load(f)


def frames2segments(frames: List[int], return_indices=False, include_length=False):
    """
    :param frames: [1, 2, 3, 10, 11, 12, 13, ...]
    :return
        (1, 3), (11, 13), ..
    """
    segments = []
    indices = []
    if len(frames) > 0:
        frame = frames[0]
        i_start = 0
        for i in range(1, len(frames)):
            if frames[i - 1] + 1 < frames[i]:
                if include_length:
                    segments.append((frame, frames[i - 1], 1 + frames[i - 1] - frame))
                else:
                    segments.append((frame, frames[i - 1]))
                indices.append((i_start, i - 1))
                frame = frames[i]
                i_start = i
        # handle the last segment
        if include_length:
            segments.append((frame, frames[-1], 1 + frames[-1] - frame))
        else:
            segments.append((frame, frames[-1]))
        indices.append((i_start, len(frames) - 1))

    if return_indices:
        assert len(indices) == len(segments)
        return segments, indices

    return segments
social_diffusion/datasets/__init__.py
jutanke/social_diffusion
import logging
import math as m
from tqdm import tqdm
from social_diffusion.transforms.rotation import rot3d
import social_diffusion.transforms.transforms as T
from social_diffusion import get_cache_dir, WeirdShapeError
from social_diffusion.sequence import (
    MultiPersonSequence,
    sequence_stats_to_unique_str,
)  # noqa E501
from social_diffusion.datasets import calculate_standardization
from social_diffusion.skeleton import LocalSkeleton
from numpy.random import uniform
from torch.utils.data import Dataset
from typing import List
import numpy as np
from einops import rearrange, reduce
from os.path import join, isfile
import hashlib
import torch


logger = logging.getLogger(__name__)


class MultiPersonDataset(Dataset):
    def __init__(
        self,
        sequences: List[MultiPersonSequence],
        n_frames: int,
        *,
        ss: int = 10,
        create_batch_entry_fn=None,
        undo_batch_fn=None,
        augmentation_fn=None,
        do_caching: bool = True,
        do_standardize: bool = True,
        p_mu=None,
        p_std=None,
        skel=None,
        calculate_standardization_fn=None,
    ):
        """
        :param create_batch_entry_fn:
            fn(P: {n_frames x n_person x jd}, skel: Skeleton) -> {batch entry}
            undo_data
        :param undo_batch_fn: fn(P: {n_frames x n_person x jd}
            kel: Skeleton, undo_data)
        :param augmentation_fn: fn(Seq: n_frames x n_person x jd}
            skel: Skeleton)
        """
        super().__init__()

        self.create_batch_entry_fn = create_batch_entry_fn
        self.undo_batch_fn = undo_batch_fn
        if augmentation_fn is None:
            self.augmentation_fn = lambda Seq, skel: Seq
        else:
            self.augmentation_fn = augmentation_fn

        if p_mu is not None and not do_standardize:
            raise ValueError(
                "You pass standardization parameter but you have standardization deactivated.. bug?"  # noqa E501
            )

        self.do_standardize = do_standardize
        if skel is None:
            self.skel = sequences[0].skel
        else:
            self.skel = skel

        self.__data = []
        n_dummy_person = sequences[0].n_max_person()
        cache_fname = None
        if do_caching:
            # --- create cache_fname ---
            cache_fname = (
                MultiPersonSequence.global_uid(sequences)
                + f"_nframes:{n_frames}_ss:{ss}"
            )
            if create_batch_entry_fn is not None:
                # we pass a dummy sequence to the preprocess_fn
                dummy = np.empty(
                    (
                        n_frames,
                        n_dummy_person,
                        self.skel.jd_without_normalization(),
                    ),  # noqa E501
                    dtype=np.float32,
                )
                incr_t = rearrange(
                    np.linspace(start=-1.0, stop=1.0 + n_frames, num=n_frames),
                    "t -> t 1 1",
                )
                incr_p = rearrange(
                    np.linspace(
                        start=-2.0, stop=n_dummy_person, num=n_dummy_person
                    ),  # noqa E501
                    "p -> 1 p 1",
                )
                for jid in range(dummy.shape[2]):
                    dummy[:, :, jid] = jid
                dummy = incr_p + dummy * incr_t
                dummy += 0.000001
                dummy_norm, _ = create_batch_entry_fn(dummy, self.skel)
                prep_str = sequence_stats_to_unique_str(
                    rearrange(dummy_norm, "t p jd -> (t p) jd")
                )
                cache_fname += prep_str

            cache_fname = join(
                get_cache_dir(),
                "mpd_"
                + hashlib.sha1(cache_fname.encode("utf-8")).hexdigest()
                + ".npy",  # noqa E501
            )

            if isfile(cache_fname):
                self.__data = np.load(cache_fname, mmap_mode="r")

        if len(self.__data) == 0:
            for seq in tqdm(
                sequences, total=len(sequences), position=0, leave=True
            ):  # noqa E501
                n_person = seq.n_max_person()
                if seq.is_masked():
                    raise ValueError("We do not allow masked sequences!")
                if n_person != n_dummy_person:
                    raise ValueError(
                        "We expect all sequences to have the same amount of persons!"  # noqa E501
                    )
                for t in range(0, len(seq) - n_frames + 1, ss):
                    P, _ = self.create_batch_entry(
                        seq.Seq[t : t + n_frames].copy()  # noqa E501
                    )  # noqa E501
                    self.__data.append(P)

            self.__data = np.array(self.__data, dtype=np.float32)
            if do_caching:
                np.save(cache_fname, self.__data)
        self.p_mu = None
        self.p_std = None
        if do_standardize:
            if p_mu is None:
                assert p_std is None

                # calculate standardization
                if calculate_standardization_fn is None:
                    logger.info("utilize default standardization function")
                    calculate_standardization_fn = calculate_standardization
                else:
                    logger.info("utilize custom standardization function")

                cache_fname_std = None
                if do_caching:
                    assert cache_fname is not None
                    cache_fname_std = cache_fname.replace(
                        ".npy", "_standard.npz"
                    )  # noqa E501
                    if isfile(cache_fname_std):
                        standard = np.load(cache_fname_std)
                        self.p_mu = standard["p_mu"]
                        self.p_std = standard["p_std"]
                if self.p_mu is None:
                    assert self.p_std is None
                    # calculate standardization
                    p_mu, p_std = calculate_standardization_fn(
                        self.__data, self.skel
                    )  # noqa E501
                    self.p_mu = rearrange(p_mu, "jd -> 1 jd").astype(
                        np.float32
                    )  # noqa E501
                    self.p_std = rearrange(p_std, "jd -> 1 jd").astype(
                        np.float32
                    )  # noqa E501
                    if do_caching:
                        np.savez(
                            cache_fname_std, p_mu=self.p_mu, p_std=self.p_std
                        )  # noqa E501
            else:
                assert p_std is not None
                self.p_mu = p_mu
                self.p_std = p_std

            # mask out standardization
            standardization_mask = skel.standardization_mask  # 1 x jd
            inverse_standardization_mask = (
                (standardization_mask < 0.5) * 1
            ).astype(  # noqa E501
                "float32"
            )
            fixed_p_std = np.ones_like(self.p_std)
            self.p_mu = self.p_mu * standardization_mask
            self.p_std = (
                self.p_std * standardization_mask
                + fixed_p_std * inverse_standardization_mask
            )

        # -- ctor over --

    def __len__(self):
        return len(self.__data)

    def __getitem__(self, index: int):
        return {
            "X_0": self.do_standardization(
                self.augmentation_fn(self.__data[index], self.skel)
            )
        }

    def get_raw(self, index: int):
        return self.__data[index]

    def undo_batch_entry(self, P, undo_data=None):
        if self.undo_batch_fn is None:
            return P
        else:
            return self.undo_batch_fn(P, self.skel, undo_data)

    def create_batch_entry(self, P):
        if self.create_batch_entry_fn is None:
            return P, None
        else:
            return self.create_batch_entry_fn(P, self.skel)

    def do_standardization(self, Seq):
        """
        :param Seq: {n_frames x 3 x jd}
        """
        if self.do_standardize:
            if not isinstance(Seq, np.ndarray):
                Seq = Seq.detach().cpu().numpy()
            assert Seq.shape[2] == self.skel.n_joints * 3
            p_mu = rearrange(self.p_mu, "b jd -> 1 b jd")
            p_std = rearrange(self.p_std, "b jd -> 1 b jd")
            return (Seq - p_mu) / p_std
        else:
            return Seq

    def undo_standardization(self, Seq):
        """
        :param Seq: {n_frames x 3 x jd}
        """
        if self.do_standardize:
            if not isinstance(Seq, np.ndarray):
                Seq = Seq.detach().cpu().numpy()
            assert Seq.shape[2] == self.skel.n_joints * 3
            p_mu = rearrange(self.p_mu, "b jd -> 1 b jd")
            p_std = rearrange(self.p_std, "b jd -> 1 b jd")
            return p_mu + (Seq * p_std)
        else:
            return Seq

    def batch_undo_standardization(self, Seq):
        """
        Keep gradients!
        :param Seq: {n_batch x n_frames x 3 x jd}
        """
        if self.do_standardize:
            p_mu = torch.from_numpy(
                rearrange(self.p_mu, "b jd -> 1 1 b jd")
            ).to(  # noqa E501
                Seq.device
            )
            p_std = torch.from_numpy(
                rearrange(self.p_std, "b jd -> 1 1 b jd")
            ).to(  # noqa E501
                Seq.device
            )
            return p_mu + (Seq * p_std)
        else:
            return Seq


def get_NormalizedMultiPersonDataset(
    sequences: List[MultiPersonSequence],
    n_frames: int,
    normalize_at_frame: int,
    *,
    ss: int = 10,
    do_caching: bool = True,
    do_standardize=True,
    # ignore_warnings=False,
) -> MultiPersonDataset:
    # JULIAN
    def create_batch_entry_fn(Seq, skel):
        if len(Seq.shape) != 3:
            raise WeirdShapeError(
                f"Weird shape for normalization: {Seq.shape}"
            )  # noqa E501

        n_person = Seq.shape[1]

        Seq_norm = []
        Mus = []
        Rs = []
        for pid in range(n_person):
            seq = Seq[:, pid]
            seq = skel.fix_shape(seq, unroll_jd=True)
            seq_norm, (mu, R) = T.normalize(
                seq,
                jid_left=skel.normalizing_left_jid,
                jid_right=skel.normalizing_right_jid,
                frame=normalize_at_frame,
                return_transform=True,
            )
            seq_norm = rearrange(seq_norm, "t j d -> t (j d)")
            Seq_norm.append(seq_norm)
            Mus.append(mu)
            Rs.append(R)

        Seq_norm = np.array(Seq_norm, dtype=np.float32)
        Seq_norm = rearrange(Seq_norm, "p t jd -> t p jd")

        return Seq_norm, {"Mus": Mus, "Rs": Rs}

    def undo_batch_fn(Seq, skel, undo_data):
        """
        :param Seq: {n_frames x n_person x jd}
        :param undo_data: {
            "Mus"
            "Rs"
        }
        """
        if len(Seq.shape) != 3:
            raise WeirdShapeError(
                f"Weird shape for normalization: {Seq.shape}"
            )  # noqa E501

        Mus = undo_data["Mus"]
        Rs = undo_data["Rs"]
        n_person = Seq.shape[1]

        Seq_recover = []
        for pid in range(n_person):
            seq = Seq[:, pid]
            Seq_recover.append(
                T.undo_normalization_to_seq(seq=seq, mu=Mus[pid], R=Rs[pid])
            )
        Seq_recover = np.array(Seq_recover, dtype=np.float32)

        # print("Seq_recover", Seq_recover.shape)
        return Seq_recover

    return MultiPersonDataset(
        sequences=sequences,
        n_frames=n_frames,
        ss=ss,
        do_caching=do_caching,
        do_standardize=do_standardize,
        # ignore_warnings=ignore_warnings,
        skel=sequences[0].skel,
        create_batch_entry_fn=create_batch_entry_fn,
        undo_batch_fn=undo_batch_fn,
    )


def get_LocalMultiPersonDataset(
    sequences: List[MultiPersonSequence],
    n_frames: int,
    normalize_at_frame: int,
    *,
    ss: int = 10,
    do_caching: bool = True,
    p_mu=None,
    p_std=None,
    do_standardize=True,
    # ignore_warnings=False,
    use_rotpoints=True,
    normalize=False,
    augmentation_random_rotate=True,
    augmentation_random_translate=False,
) -> MultiPersonDataset:
    def create_batch_entry_fn(Seq, skel):
        if len(Seq.shape) != 3:
            raise WeirdShapeError(
                f"Weird shape for normalization: {Seq.shape}"
            )  # noqa E501

        n_person = Seq.shape[1]

        Seq_norm = []
        for pid in range(n_person):
            seq = Seq[:, pid]
            seq = skel.to_local(seq, use_rotpoints=use_rotpoints, is_flat=True)
            Seq_norm.append(seq)
        Seq_norm = np.array(Seq_norm, dtype=np.float32)
        Seq_norm = rearrange(Seq_norm, "p t jd -> t p jd")
        # --- v ---
        # print("== WARNING == WE DESTROY POSE + ROTATION HERE!")
        # t, p, jd = Seq_norm.shape
        # Seq_norm[:, :, 3:] = np.random.uniform(
        #     low=-0.01, high=0.01, size=(t, p, jd - 3)
        # )
        # print("== WARNING == WE DESTROY POSE + ROTATION HERE!")
        # --- ^ ---

        if normalize:
            # Last_loc_for_all = reduce(
            #     Seq_norm[normalize_at_frame - 1, :, 0:3], "p d -> d", "mean"
            # )
            # Last_loc_for_all[2] = 0
            # t = rearrange(Last_loc_for_all, "d -> 1 1 d")
            t = rearrange(Seq_norm[normalize_at_frame, :, :3], "p d -> 1 p d")
            # Seq_norm[:, :, :3] = Seq_norm[:, :, :3] - t
            # return Seq_norm, {"translate": t}
            Seq_norm[:, :, :3] = Seq_norm[:, :, :3] - t
            return Seq_norm, {"translate": t}
        else:
            return Seq_norm, None

    def undo_batch_fn(Seq, skel, undo_data):
        """
        :param Seq: {n_frames x n_person x jd}
        :param undo_data: {
            "translate", 1x3 for each n_person
        }
        """
        if len(Seq.shape) != 3:
            raise WeirdShapeError(
                f"Weird shape for normalization: {Seq.shape}"
            )  # noqa E501

        if undo_data is not None:
            t = undo_data["translate"]
            Seq[:, :, :3] += t

        return skel.from_local_group(Seq, use_rotpoints=use_rotpoints)

    def augmentation_fn(Seq, skel):
        if (
            not augmentation_random_rotate
            and not augmentation_random_translate  # noqa E501
        ):  # noqa E501
            return Seq
        else:
            if augmentation_random_rotate:
                if len(Seq.shape) != 3:
                    raise WeirdShapeError(f"Weird shape: {Seq.shape}")
                angle = uniform(low=0.0, high=m.pi * 2)
            else:
                angle = 0

            R = rot3d(0.0, 0.0, angle)

            if augmentation_random_translate:
                t = uniform(low=-3.0, high=3.0, size=(1, 1, 1, 3))
                t[:, :, :, 2] = 0
            else:
                t = np.zeros((1, 1, 1, 3), dtype=np.float32)

            n_person = Seq.shape[1]
            Seq_out = []
            for pid in range(n_person):
                seq = skel.fix_shape(
                    Seq[:, pid].copy(), unroll_jd=True
                )  # n_frames x j x 3
                seq[:, 0] = seq[:, 0] @ R + t
                seq[:, 1] = seq[:, 1] @ R  # WE ASSUME THIS IS rotpoint!
                Seq_out.append(seq)
            return rearrange(Seq_out, "p t j d -> t p (j d)")

    calculate_standardization_fn = None
    if augmentation_random_rotate or augmentation_random_translate:
        # if we randomly rotate (augment) we HAVE to adjust the standardization
        # function!

        def calculate_standardization_fn(all_poses, skel):
            """
            :param all_poses: {b t p jd}
            """
            eps = 0.000000001
            assert len(all_poses.shape) == 4, f"weird shape: {all_poses.shape}"
            assert skel.is_local
            all_poses = rearrange(
                all_poses, "b t p (j d) -> b t p j d", j=skel.n_joints
            )

            loc = rearrange(all_poses[:, :, :, 0], "b t p d -> (b t p) d")
            rot = rearrange(all_poses[:, :, :, 1], "b t p d -> (b t p) d")
            pose = rearrange(
                all_poses[:, :, :, 2:], "b t p j d -> (b t p) (j d)"
            )  # noqa E501

            Locs = []
            Rots = []
            if augmentation_random_rotate:
                for angle in np.linspace(start=0, stop=2 * m.pi, num=32):
                    R = rot3d(0.0, 0.0, angle)
                    Locs.append(loc.copy() @ R)
                    Rots.append(rot.copy() @ R)
            else:
                Rots.append(rot)

            if augmentation_random_translate:
                for x in np.linspace(start=-3, stop=3, num=10):
                    for y in np.linspace(start=-3, stop=3, num=10):
                        t = np.array([x, y, 0], dtype=np.float32)
                        Locs.append(loc.copy() + t)
            else:
                Locs.append(loc)

            Locs = np.concatenate(Locs, axis=0)
            Rots = np.concatenate(Rots, axis=0)

            p_mu_pose = reduce(pose, "t jd -> jd", "mean")
            p_std_pose = np.std(pose, axis=0) + eps

            p_mu_loc = reduce(Locs, "t jd -> jd", "mean")
            p_std_loc = np.std(Locs, axis=0) + eps

            p_mu_rot = reduce(Rots, "t jd -> jd", "mean")
            p_std_rot = np.std(Rots, axis=0) + eps

            p_mu = np.concatenate([p_mu_loc, p_mu_rot, p_mu_pose], axis=0)
            p_std = np.concatenate([p_std_loc, p_std_rot, p_std_pose], axis=0)

            return p_mu, p_std

    parent_skel = sequences[0].skel
    skel = LocalSkeleton(parent_skel)
    return MultiPersonDataset(
        sequences=sequences,
        n_frames=n_frames,
        ss=ss,
        do_caching=do_caching,
        p_mu=p_mu,
        p_std=p_std,
        do_standardize=do_standardize,
        # ignore_warnings=ignore_warnings,
        skel=skel,
        create_batch_entry_fn=create_batch_entry_fn,
        undo_batch_fn=undo_batch_fn,
        augmentation_fn=augmentation_fn,
        calculate_standardization_fn=calculate_standardization_fn,
    )


def get_LocalOnlyMultiPersonDataset(
    sequences: List[MultiPersonSequence],
    n_frames: int,
    ss: int = 10,
    do_caching: bool = True,
    p_mu=None,
    p_std=None,
    do_standardize=True,
    # ignore_warnings=False,
):
    """
    This is a **debugging** module! Only use local rep!
    """

    def create_batch_entry_fn(Seq, skel):
        if len(Seq.shape) != 3:
            raise WeirdShapeError(
                f"Weird shape for normalization: {Seq.shape}"
            )  # noqa E501

        n_person = Seq.shape[1]

        Seq_norm = []
        for pid in range(n_person):
            seq = Seq[:, pid]
            seq = skel.to_local(seq, use_rotpoints=True, is_flat=True)
            Seq_norm.append(seq)
        Seq_norm = np.array(Seq_norm, dtype=np.float32)
        Seq_norm = rearrange(Seq_norm, "p t (j d) -> t p j d", d=3)
        Seq_norm = rearrange(Seq_norm[:, :, 2:], "t p j d -> t p (j d)")
        return Seq_norm, None

    skel = sequences[0].skel
    return MultiPersonDataset(
        sequences=sequences,
        n_frames=n_frames,
        ss=ss,
        do_caching=do_caching,
        p_mu=p_mu,
        p_std=p_std,
        do_standardize=do_standardize,
        # ignore_warnings=ignore_warnings,
        skel=skel,
        create_batch_entry_fn=create_batch_entry_fn,
    )
social_diffusion/datasets/dataset.py
jutanke/social_diffusion
import logging
import pickle
import numpy as np
import numpy.linalg as la
import numba as nb
from social_diffusion.skeleton import Skeleton
from social_diffusion.sequence import MultiPersonSequence
from social_diffusion.datasets.panoptic import (
    PANOPTIC_SKELETON,
    interpolate_missing_limbs,
)
from social_diffusion.datasets import frames2segments

from social_diffusion.datasets.dataset import get_LocalMultiPersonDataset
from social_diffusion import get_data_dir
from multiprocessing.pool import ThreadPool
from einops import rearrange, reduce
import torch
from scipy.ndimage import gaussian_filter1d
from os.path import join
import math as m
from typing import List
from time import time
from tqdm import tqdm


logger = logging.getLogger(__name__)


def get_haggling_local_dataset(n_frames: int, n_in: int):
    train_seq = get_haggling_train_sequences()
    test_seq = get_haggling_test_sequences()
    ds = get_LocalMultiPersonDataset(  # noqa E501
        train_seq,
        n_frames=n_frames,
        normalize_at_frame=n_in - 1,
        normalize=False,
        augmentation_random_rotate=False,
        augmentation_random_translate=False,
    )
    ds_test = get_LocalMultiPersonDataset(  # noqa E501
        test_seq,
        n_frames=n_frames,
        normalize_at_frame=n_in - 1,
        normalize=False,
        augmentation_random_rotate=False,
        augmentation_random_translate=False,
    )
    return ds, ds_test, ds.skel


def get_haggling_test_sequences(to_sequence=True) -> List[MultiPersonSequence]:
    """ """
    _start = time()
    TEST = HagglingScene.load_test()
    logger.info(
        "get_haggling_test_sequences => elapsed :%0.4f" % (time() - _start)
    )  # noqa E501
    for scene in tqdm(TEST):
        scene = manually_fix(scene)
        scene.fix_poses()
    if to_sequence:
        return [scene.to_sequence() for scene in TEST]
    else:
        return [scene for scene in TEST]


def get_haggling_train_sequences(to_sequence=True) -> List[MultiPersonSequence]:
    """ """
    _start = time()
    TRAIN = HagglingScene.load_train()
    logger.info(
        "get_haggling_train_sequences => elapsed :%0.4f" % (time() - _start)
    )  # noqa E501
    for scene in tqdm(TRAIN):
        scene = manually_fix(scene)
        scene.fix_poses()
    if to_sequence:
        return [scene.to_sequence() for scene in TRAIN]
    else:
        return [scene for scene in TRAIN]
    # return [scene.to_sequence() for scene in TRAIN]


@nb.njit(nb.float32[:, :, :, :](nb.float32[:, :, :, :]), nogil=True)
def backfill_zeros(seq):
    """
    IF the method starts with zero entries we have to find the first good
    entry!
    :param seq: {n_frames x n_person x J x 4}
    """
    n_frames = seq.shape[0]
    n_person = seq.shape[1]
    n_joints = seq.shape[2]
    for pid in range(n_person):
        for jid in range(n_joints):
            if seq[0, pid, jid, 3] < 0.09:
                for t in range(1, n_frames):
                    if seq[t, pid, jid, 3] > 0.15:
                        for tt in range(0, t):
                            seq[tt, pid, jid, :3] = seq[t, pid, jid, :3]
    return seq


class HagglingScene:
    @staticmethod
    def skeleton() -> Skeleton:
        return PANOPTIC_SKELETON

    @staticmethod
    def load_all():
        global ALL_FILES_NAMES
        Scenes = []
        for scene_name in ALL_FILES_NAMES:
            scene = HagglingScene(scene_name)
            Scenes.append(scene)
        return Scenes

    @staticmethod
    def load_train():
        global ALL_FILES_NAMES
        TRAIN_FILES = [
            scene_name
            for scene_name in ALL_FILES_NAMES
            if scene_name[: scene_name.find("_group")] not in TESTING_SEQUENCES
        ]
        with ThreadPool(len(TRAIN_FILES)) as p:
            return p.map(HagglingScene, TRAIN_FILES)

    @staticmethod
    def load_test():
        global ALL_FILES_NAMES
        Scenes = []
        for scene_name in ALL_FILES_NAMES:
            full_name = scene_name[: scene_name.find("_group")]
            if full_name in TESTING_SEQUENCES:
                scene = HagglingScene(scene_name)
                Scenes.append(scene)
        return Scenes

    def __init__(self, scene_name: str):
        if not scene_name.endswith(".pkl"):
            raise ValueError(f"Name should end on .pkl ({scene_name})")
        self.scene_name = scene_name

        local_path = join(get_data_dir(), "haggling")

        fname_pose = join(
            local_path, f"panopticDB_body_pkl_hagglingProcessed/{scene_name}"
        )  # noqa E501
        fname_speech = join(
            local_path, f"panopticDB_speech_pkl_hagglingProcessed/{scene_name}"
        )  # noqa E501

        with open(fname_pose, "rb") as f:
            P = pickle.load(f, encoding="latin1")
            self.seller_ids = list(sorted(P["sellerIds"]))
            self.left_seller_id = P["leftSellerId"]
            self.right_seller_id = P["rightSellerId"]
            self.buyer_id = P["buyerId"]
            self.winner_id = P["winnerId"]
            self.poses = {}  # pid -> seq
            self.scores = {}  # pid -> scores
            self.normals = {}  # pid -> {"body", "face"}
            for i in range(3):
                J = rearrange(P["subjects"][i]["joints19"], "d t -> t d")
                body_normal = rearrange(
                    P["subjects"][i]["bodyNormal"], "d t -> t d"
                )  # noqa E501
                face_normal = rearrange(
                    P["subjects"][i]["faceNormal"], "d t -> t d"
                )  # noqa E501
                scores = rearrange(P["subjects"][i]["scores"], "d t -> t d")
                # make z point upwards - - -
                z_flip = np.diag([1, 1, -1])
                J = rearrange(J, "t (j d) -> (t j) d", j=19, d=3)
                J[:, [0, 1, 2]] = (J[:, [0, 2, 1]] / 100) @ z_flip
                J = rearrange(J, "(t j) d -> t (j d)", j=19, d=3)

                # we have to flip body/face normals as well!
                body_normal[:, [0, 1, 2]] = (
                    body_normal[:, [0, 2, 1]]
                ) @ z_flip  # noqa E501
                face_normal[:, [0, 1, 2]] = (
                    face_normal[:, [0, 2, 1]]
                ) @ z_flip  # noqa E501
                # - - -

                pid = P["subjects"][i]["humanId"]
                self.poses[pid] = J
                self.scores[pid] = scores

        with open(fname_speech, "rb") as f:
            S = pickle.load(f, encoding="latin1")
            self.speech = {}  # pid -> speech
            self.words = {}  # pid -> words
            for i in range(3):
                pid = S["speechData"][i]["humanId"]
                speech = S["speechData"][i]["indicator"][
                    :-1
                ]  # speech has one entry too many
                self.speech[pid] = speech

        self.n_frames = -1
        for pid in self.poses.keys():
            # assert pid in self.speech
            if pid not in self.speech:
                raise ValueError(f"Cannot find pid:{pid} in speech!")
            if self.n_frames == -1:
                self.n_frames = len(self.poses[pid])
            else:
                if self.n_frames != len(self.poses[pid]):
                    raise ValueError(
                        f"frames and poses inconsistent: {self.n_frames} vs {len(self.poses[pid])}"  # noqa E501
                    )  # noqa E501
            if len(self.speech[pid]) != self.n_frames:
                raise ValueError(
                    f"pid{pid} speech:{self.speech[pid].shape}, poses:{self.poses[pid].shape}"  # noqa E501
                )  # noqa E501

        self.pids = sorted(self.poses.keys())
        self.pid2index = {}
        for index, pid in enumerate(self.pids):
            self.pid2index[pid] = index

    def __len__(self):
        return self.n_frames

    def to_sequence(self) -> MultiPersonSequence:
        return MultiPersonSequence(
            Seq=self.get_poses(), skel=HagglingScene.skeleton()
        )  # noqa E501

    def mean_location(self):
        mu = reduce(
            self.get_poses(), "t p (j d) -> d", reduction="mean", j=19, d=3
        )  # noqa E501
        mu[2] = 0
        return mu

    def which_seller_has_the_attention(
        self, return_speech=False, return_buyer_speech=False
    ):
        if return_buyer_speech:
            if not return_speech:
                raise ValueError("which_seller_has_the_attention fail")
        HO = self.calculate_head_orientations()

        pid2idx = {}
        for i, pid in enumerate(self.pids):
            pid2idx[pid] = i

        seller1 = HO[:, pid2idx[self.seller_ids[0]]]
        seller2 = HO[:, pid2idx[self.seller_ids[1]]]
        buyer = HO[:, pid2idx[self.buyer_id]]

        seller1_loc = seller1[:, 0]
        seller2_loc = seller2[:, 0]
        buyer_loc = buyer[:, 0]
        buyer_lookdir = buyer[:, 1]

        def get_dir_vec(a, b):
            ab = b - a
            return ab / la.norm(ab, axis=1, keepdims=True)

        b_to_s1 = get_dir_vec(buyer_loc, seller1_loc)
        b_to_s2 = get_dir_vec(buyer_loc, seller2_loc)

        att1 = np.einsum("ij,ij->i", b_to_s1, buyer_lookdir)
        att2 = np.einsum("ij,ij->i", b_to_s2, buyer_lookdir)

        attn = np.zeros((self.n_frames, 2), dtype=np.int64)
        for t in range(self.n_frames):
            if att1[t] > att2[t]:
                attn[t, 0] = 1
            else:
                attn[t, 1] = 1

        if return_speech:
            all_speech = self.get_speech()  # n_frames x n_person
            speech1 = rearrange(
                all_speech[:, pid2idx[self.seller_ids[0]]], "t -> t 1"
            )  # noqa E501
            speech2 = rearrange(
                all_speech[:, pid2idx[self.seller_ids[1]]], "t -> t 1"
            )  # noqa E501
            speech = np.concatenate([speech1, speech2], axis=1)

            if return_buyer_speech:
                buyer_speech = all_speech[:, pid2idx[self.buyer_id]]
                return attn, speech, buyer_speech
            else:
                return attn, speech
        else:
            return attn

    def resize(self, start: int, end: int):
        for pid in self.pids:
            self.poses[pid] = self.poses[pid][start:end]
            self.speech[pid] = self.speech[pid][start:end]
            self.scores[pid] = self.scores[pid][start:end]
        self.n_frames = len(self.poses[pid])

    def fix_poses(self):
        n_person = 3
        m = np.ones((self.n_frames), dtype=np.int64)
        Seq = self.get_poses_with_scores()

        # step 1: backfill
        # Seq = backfill_zeros(Seq)

        with ThreadPool(n_person) as p:
            PM_new = p.starmap(
                interpolate_missing_limbs,
                [(Seq[:, i], m, i) for i in range(n_person)],  # noqa E501
            )
            P_new = np.array([P for P, _ in PM_new], dtype=np.float32)

        for i, pid in enumerate(self.pids):
            seq = P_new[i]  # t x 19 x 4
            seq = seq[:, :, :3]
            self.poses[pid] = rearrange(seq, "t j d -> t (j d)")

    def get_poses_with_scores(self, sigma=-1.0):
        poses = rearrange(
            self.get_poses(sigma=sigma), "t p (j d) -> t p j d", j=19, d=3
        )
        scores = rearrange(self.get_scores(), "t p j -> t p j 1")
        return np.concatenate([poses, scores], axis=3)

    def get_scores(self):
        scores = np.empty((self.n_frames, 3, 19), dtype=np.float32)
        for i, pid in enumerate(self.pids):
            scores[:, i] = self.scores[pid]
        return scores

    def get_poses(self, sigma=0.5):
        poses = np.empty((self.n_frames, 3, 57), dtype=np.float32)
        for i, pid in enumerate(self.pids):
            poses[:, i] = self.poses[pid]
        if sigma > 0:
            poses = smooth_poses(poses, sigma=sigma)
        return poses

    def get_ordered_poses(self, sigma=0.5):
        P = self.get_poses(sigma=sigma)

        buyer_index = self.pid2index[self.buyer_id]
        left_index = self.pid2index[self.left_seller_id]
        right_index = self.pid2index[self.right_seller_id]

        return np.ascontiguousarray(
            P[:, [buyer_index, left_index, right_index]]
        )  # noqa E501

    def update_poses(self, poses: np.ndarray, scores: np.ndarray = None):
        assert len(poses) == self.n_frames
        for idx in range(3):
            pid = self.pids[idx]
            self.poses[pid] = poses[:, idx]
        if scores is not None:
            assert len(scores) == self.n_frames
            for idx in range(3):
                pid = self.pids[idx]
                self.scores[pid] = scores[:, idx]

    def get_plot_description(self):
        """
        return:
            per_person_text: {
                pid, t --> "buyer|<talk>"
            }
        """
        per_person_text = {}
        for t in range(self.n_frames):
            for i, pid in enumerate(self.pids):
                txt = ""
                if pid == self.buyer_id:
                    txt += "<buyer>"
                elif pid == self.left_seller_id:
                    txt += "<left seller>"
                elif pid == self.right_seller_id:
                    txt += "<right seller>"
                else:
                    raise ValueError(
                        f"Cannot assign {pid} for {self.scene_name}"
                    )  # noqa E501
                if self.speech[pid][t] > 0.5:
                    txt += "|<speak>"
                per_person_text[i, t] = txt
        return per_person_text

    def get_speech(self):
        speech = np.empty((self.n_frames, 3), dtype=np.int64)
        for i, pid in enumerate(self.pids):
            speech[:, i] = self.speech[pid]
        return speech

    def calculate_head_orientations(self):
        """
        Our OWN head orientation calculation
        return:
            {n_frames x n_person x 2 x 2}  # (pos / norm)
        """
        n_person = 3
        poses = self.get_poses()
        head_orientations = np.empty(
            (self.n_frames, n_person, 2, 2), dtype=np.float32
        )  # noqa E501
        for i in range(n_person):
            seq = poses[:, i]
            head_orientations[:, i, 0] = rearrange(
                seq, "t (j d) -> t j d", j=19, d=3
            )[  # noqa E501
                :, 1, :2
            ]
            head_orientations[:, i, 1] = get_normal(
                seq, left_jid=15, right_jid=17
            )  # noqa E501
        return head_orientations

    def calculate_body_orientations(self):
        n_person = 3
        poses = self.get_poses()
        body_orientations = np.empty(
            (self.n_frames, n_person, 2, 2), dtype=np.float32
        )  # noqa E501
        for i in range(n_person):
            seq = poses[:, i]
            body_orientations[:, i, 0] = rearrange(
                seq, "t (j d) -> t j d", j=19, d=3
            )[  # noqa E501
                :, 2, :2
            ]
            body_orientations[:, i, 1] = get_normal(
                seq, left_jid=6, right_jid=12
            )  # noqa E501
        return body_orientations

    def has_problematic_lower_parts(self):
        txt = "- - - - - - - - - - - - - - - - - - "
        P = self.get_poses_with_scores()
        something_is_broken = False
        for pid in range(3):
            pp = P[:, pid, [6, 7, 8, 12, 13, 14], 3]
            broken_lp_indices = ((np.min(pp, axis=1) < 0) * 1).nonzero()[0]
            if len(broken_lp_indices) > 0:
                for start, end, length in frames2segments(
                    broken_lp_indices, include_length=True
                ):
                    if length > 10:
                        something_is_broken = True
                        txt += f"\npid{pid} -> {start} to {end} @{self.scene_name}"  # noqa E501
        if something_is_broken:
            print(txt)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# U T I L S
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


def rot2d(theta):
    return np.array(
        [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]],
        dtype=np.float32,
    )


def get_normal(poses, left_jid: int, right_jid: int, smooth=True):
    """
    :param poses: {n_frames x 57}
    """
    is_torch = True
    if isinstance(poses, np.ndarray):
        is_torch = False
        poses = torch.from_numpy(poses)
    poses = rearrange(poses, "t (j d) -> t j d", j=19, d=3)
    left2d = poses[:, left_jid, :2]
    right2d = poses[:, right_jid, :2]
    y = right2d - left2d  # n_frames x 2
    y = y / (torch.norm(y, dim=1, p=2).unsqueeze(1) + 0.00000001)
    if not is_torch:
        y = y.numpy()
    # y[:, [0, 1]] = y[:, [1, 0]]

    R = rot2d(-m.pi / 2)

    y = y @ R

    if smooth:
        sigma = 5.0
        y[:, 0] = gaussian_filter1d(y[:, 0], sigma=sigma)
        y[:, 1] = gaussian_filter1d(y[:, 1], sigma=sigma)

    y = y / la.norm(y, axis=1, keepdims=True)

    return y


def smooth_poses(poses, sigma, sigma_hip=1.5):
    """
    :param poses: {n_frames x n_person x 56}
    """
    hip_dims = set(
        [
            2 * 3,
            2 * 3 + 1,
            2 * 3 + 2,
            6 * 3,
            6 * 3 + 1,
            6 * 3 + 2,
            12 * 3,
            12 * 3 + 1,
            12 * 3 + 2,
        ]
    )

    n_person = poses.shape[1]
    n_dim = poses.shape[2]
    assert len(poses.shape) == 3
    for pid in range(n_person):
        for dim in range(n_dim):
            if dim in hip_dims:
                _sigma = sigma_hip
            else:
                _sigma = sigma
            poses[:, pid, dim] = gaussian_filter1d(
                poses[:, pid, dim], sigma=_sigma
            )  # noqa E501
    return poses


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ===== M A N U A L L Y  F I X =====
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
TESTING_SEQUENCES = set(
    [  # https://github.com/CMU-Perceptual-Computing-Lab/ssp
        "170221_haggling_b1",
        "170221_haggling_b2",
        "170221_haggling_b3",
        "170228_haggling_b1",
        "170228_haggling_b2",
        "170228_haggling_b3",
    ]
)

ALL_FILES_NAMES = [
    "170221_haggling_b1_group0.pkl",
    "170221_haggling_b1_group1.pkl",
    "170221_haggling_b1_group2.pkl",
    "170221_haggling_b1_group3.pkl",
    "170221_haggling_b1_group4.pkl",
    "170221_haggling_b1_group5.pkl",
    "170221_haggling_b2_group0.pkl",
    "170221_haggling_b2_group1.pkl",
    "170221_haggling_b2_group2.pkl",
    "170221_haggling_b2_group3.pkl",
    "170221_haggling_b2_group4.pkl",
    "170221_haggling_b2_group5.pkl",
    "170221_haggling_b3_group0.pkl",
    "170221_haggling_b3_group1.pkl",
    "170221_haggling_b3_group2.pkl",
    "170221_haggling_m1_group0.pkl",
    "170221_haggling_m1_group1.pkl",
    "170221_haggling_m1_group2.pkl",
    "170221_haggling_m1_group3.pkl",
    "170221_haggling_m1_group4.pkl",
    "170221_haggling_m1_group5.pkl",
    "170221_haggling_m2_group0.pkl",
    "170221_haggling_m2_group1.pkl",
    "170221_haggling_m2_group2.pkl",
    "170221_haggling_m2_group3.pkl",
    "170221_haggling_m2_group4.pkl",
    "170221_haggling_m2_group5.pkl",
    "170221_haggling_m3_group0.pkl",
    "170221_haggling_m3_group1.pkl",
    "170221_haggling_m3_group2.pkl",
    "170224_haggling_a1_group0.pkl",
    "170224_haggling_a1_group1.pkl",
    # "170224_haggling_a1_group2.pkl",  # broken lower body parts
    "170224_haggling_a1_group3.pkl",
    "170224_haggling_a1_group4.pkl",
    "170224_haggling_a1_group5.pkl",
    "170224_haggling_a1_group6.pkl",
    # "170224_haggling_a1_group7.pkl",  # broken lower body parts
    "170224_haggling_a2_group0.pkl",
    "170224_haggling_a2_group1.pkl",
    "170224_haggling_a2_group2.pkl",
    "170224_haggling_a2_group3.pkl",
    # "170224_haggling_a2_group4.pkl",  # broken lower body parts
    "170224_haggling_a2_group5.pkl",
    "170224_haggling_a2_group6.pkl",
    "170224_haggling_a2_group7.pkl",
    "170224_haggling_a3_group0.pkl",
    # "170224_haggling_a3_group1.pkl",  # broken lower body parts
    # "170224_haggling_a3_group2.pkl",  # broken lower body parts
    "170224_haggling_a3_group3.pkl",
    "170224_haggling_b1_group0.pkl",
    "170224_haggling_b1_group1.pkl",
    "170224_haggling_b1_group2.pkl",
    "170224_haggling_b1_group3.pkl",
    "170224_haggling_b1_group4.pkl",
    "170224_haggling_b1_group5.pkl",
    "170224_haggling_b1_group6.pkl",
    "170224_haggling_b1_group7.pkl",
    "170224_haggling_b2_group0.pkl",
    "170224_haggling_b2_group1.pkl",
    "170224_haggling_b2_group2.pkl",
    "170224_haggling_b2_group3.pkl",
    "170224_haggling_b2_group4.pkl",
    "170224_haggling_b2_group5.pkl",
    "170224_haggling_b2_group6.pkl",
    "170224_haggling_b2_group7.pkl",
    "170224_haggling_b3_group0.pkl",
    "170224_haggling_b3_group1.pkl",
    "170224_haggling_b3_group2.pkl",
    "170224_haggling_b3_group3.pkl",
    "170228_haggling_a1_group0.pkl",
    "170228_haggling_a1_group1.pkl",
    "170228_haggling_a1_group2.pkl",
    "170228_haggling_a1_group3.pkl",
    "170228_haggling_a1_group4.pkl",
    "170228_haggling_a1_group5.pkl",
    "170228_haggling_a1_group6.pkl",
    "170228_haggling_a1_group7.pkl",
    "170228_haggling_a2_group0.pkl",
    "170228_haggling_a2_group1.pkl",
    "170228_haggling_a2_group2.pkl",
    # "170228_haggling_a2_group3.pkl",  # BROKEN
    "170228_haggling_a2_group4.pkl",
    "170228_haggling_a2_group5.pkl",
    "170228_haggling_a2_group6.pkl",
    "170228_haggling_a2_group7.pkl",
    "170228_haggling_a3_group0.pkl",
    "170228_haggling_a3_group1.pkl",
    "170228_haggling_a3_group2.pkl",
    "170228_haggling_a3_group3.pkl",
    "170228_haggling_b1_group0.pkl",
    "170228_haggling_b1_group1.pkl",
    "170228_haggling_b1_group2.pkl",
    "170228_haggling_b1_group3.pkl",
    "170228_haggling_b1_group4.pkl",
    "170228_haggling_b1_group5.pkl",
    "170228_haggling_b1_group6.pkl",
    "170228_haggling_b1_group7.pkl",
    "170228_haggling_b1_group8.pkl",
    "170228_haggling_b1_group9.pkl",
    "170228_haggling_b2_group0.pkl",
    "170228_haggling_b2_group1.pkl",
    "170228_haggling_b2_group2.pkl",
    "170228_haggling_b2_group3.pkl",
    "170228_haggling_b2_group4.pkl",
    "170228_haggling_b2_group5.pkl",
    "170228_haggling_b2_group6.pkl",
    "170228_haggling_b2_group7.pkl",
    "170228_haggling_b2_group8.pkl",
    "170228_haggling_b2_group9.pkl",
    "170228_haggling_b3_group0.pkl",
    "170228_haggling_b3_group1.pkl",
    "170228_haggling_b3_group2.pkl",
    "170228_haggling_b3_group3.pkl",
    "170228_haggling_b3_group4.pkl",
    "170404_haggling_a1_group0.pkl",
    "170404_haggling_a1_group1.pkl",
    "170404_haggling_a1_group2.pkl",
    "170404_haggling_a1_group3.pkl",
    "170404_haggling_a2_group0.pkl",
    "170404_haggling_a2_group1.pkl",
    "170404_haggling_a2_group2.pkl",
    "170404_haggling_a2_group3.pkl",
    "170404_haggling_a3_group0.pkl",
    "170404_haggling_a3_group1.pkl",
    "170404_haggling_b1_group0.pkl",
    "170404_haggling_b1_group1.pkl",
    "170404_haggling_b1_group2.pkl",
    "170404_haggling_b1_group3.pkl",
    "170404_haggling_b1_group4.pkl",
    "170404_haggling_b1_group5.pkl",
    "170404_haggling_b1_group6.pkl",
    "170404_haggling_b1_group7.pkl",
    "170404_haggling_b2_group0.pkl",
    "170404_haggling_b2_group1.pkl",
    "170404_haggling_b2_group2.pkl",
    "170404_haggling_b2_group3.pkl",
    "170404_haggling_b2_group4.pkl",
    "170404_haggling_b2_group5.pkl",
    "170404_haggling_b2_group6.pkl",
    "170404_haggling_b2_group7.pkl",
    "170404_haggling_b3_group0.pkl",
    "170404_haggling_b3_group1.pkl",
    "170404_haggling_b3_group2.pkl",
    "170404_haggling_b3_group3.pkl",
    # "170407_haggling_a1_group0.pkl",  # BROKEN
    "170407_haggling_a1_group1.pkl",
    "170407_haggling_a1_group2.pkl",
    "170407_haggling_a1_group3.pkl",
    "170407_haggling_a1_group4.pkl",
    "170407_haggling_a1_group5.pkl",
    "170407_haggling_a2_group0.pkl",
    "170407_haggling_a2_group1.pkl",
    "170407_haggling_a2_group2.pkl",
    "170407_haggling_a2_group3.pkl",
    "170407_haggling_a2_group4.pkl",
    "170407_haggling_a2_group5.pkl",
    "170407_haggling_a3_group0.pkl",
    "170407_haggling_a3_group1.pkl",
    "170407_haggling_a3_group2.pkl",
    "170407_haggling_b1_group0.pkl",
    "170407_haggling_b1_group1.pkl",
    "170407_haggling_b1_group2.pkl",
    "170407_haggling_b1_group3.pkl",
    "170407_haggling_b1_group4.pkl",
    "170407_haggling_b1_group5.pkl",
    "170407_haggling_b1_group6.pkl",
    "170407_haggling_b1_group7.pkl",
    "170407_haggling_b2_group0.pkl",
    "170407_haggling_b2_group1.pkl",
    "170407_haggling_b2_group2.pkl",
    "170407_haggling_b2_group3.pkl",
    "170407_haggling_b2_group4.pkl",
    "170407_haggling_b2_group5.pkl",
    "170407_haggling_b2_group6.pkl",
    "170407_haggling_b2_group7.pkl",
    "170407_haggling_b3_group0.pkl",
    "170407_haggling_b3_group1.pkl",
    "170407_haggling_b3_group2.pkl",
    "170407_haggling_b3_group3.pkl",
]


def manually_fix(scene: HagglingScene):
    """
    :param sequence:  {Scene}
    """
    global MANUAL_FIX_LEGS, MANUAL_FIX, MANUAL_RESIZE

    # --- GENERAL BODY ---
    if scene.scene_name in MANUAL_FIX:
        for entry in MANUAL_FIX[scene.scene_name]:
            pid = entry["pid"]  # not actual pid but index into poses
            start = entry["start"]
            end = entry["end"]
            jids = entry["jids"]

            poses = rearrange(
                scene.get_poses(), "t p (j d) -> t p j d", j=19, d=3
            )  # noqa E501
            scores = scene.get_scores()  # t p j
            for jid in jids:
                left = poses[start, pid, jid]
                right = poses[end, pid, jid]
                interp = np.linspace(left, right, num=end - start)
                for idx, t in enumerate(range(start + 1, end)):
                    tar = interp[idx]
                    cur = poses[t, pid, jid]
                    if la.norm(cur - tar) > 0.1:
                        poses[t, pid, jid] = tar
                        scores[t, pid, jid] = 2.0
            scene.update_poses(
                rearrange(poses, "t p j d -> t p (j d)"), scores
            )  # noqa E501

    # --- LEGS ---
    LEG_JIDS = [2, 6, 7, 8, 12, 13, 14]
    LEG_JID_INDICES = []
    for jid in LEG_JIDS:
        LEG_JID_INDICES.append(jid * 3)
        LEG_JID_INDICES.append(jid * 3 + 1)
        LEG_JID_INDICES.append(jid * 3 + 2)

    if scene.scene_name in MANUAL_FIX_LEGS:
        poses = scene.get_poses()  # n_frames x n_person x 57
        scores = scene.get_scores()  # t p j
        for entry in MANUAL_FIX_LEGS[scene.scene_name]:
            pid = entry["pid"]  # not actual pid but index into poses
            good_frames = entry["good_frames"]

            # step 0 - pre-fill to 0
            # TODO

            # step 1 - fill-up in the middle
            for i in range(len(good_frames) - 1):
                frame_a = good_frames[i]
                frame_b = good_frames[i + 1]
                if frame_a >= frame_b:  # assert frame_a < frame_b
                    raise ValueError(f"{frame_a} vs {frame_b}")
                pose_a = poses[frame_a, pid]
                pose_b = poses[frame_b, pid]
                interp = np.linspace(pose_a, pose_b, num=frame_b - frame_a)
                for j, t in enumerate(range(frame_a + 1, frame_b)):
                    for idx in LEG_JID_INDICES:
                        tar = interp[j, idx]
                        cur = poses[t, pid, idx]
                        if abs(cur - tar) > 0.1:
                            poses[t, pid, idx] = tar
                            scores[t, pid, jid] = 2.0

            # step 3 - fill-up to end
            # TODO

            # step 4 - make sure that <2> is ~ btwn <6> & <12>
            poses = rearrange(poses, "t p (j d) -> t p j d", j=19, d=3)
            n_frames = len(poses)
            for t in range(0, n_frames):
                center = poses[t, pid, 2]
                left = poses[t, pid, 6]
                right = poses[t, pid, 12]
                supposed_center = (left + right) / 2
                d = la.norm(supposed_center - center)
                if d > 0.01:
                    poses[t, pid, 2] = supposed_center
                    scores[t, pid, 2] = 2.0
            poses = rearrange(poses, "t p j d -> t p (j d)")

        # scene.poses = poses
        scene.update_poses(poses, scores)

    if scene.scene_name in MANUAL_RESIZE:
        entry = MANUAL_RESIZE[scene.scene_name]
        start = 0
        end = len(scene)
        if "start" in entry:
            start = entry["start"]
        if "end" in entry:
            end = entry["end"]
        scene.resize(start, end)

    return scene


MANUAL_RESIZE = {
    "170228_haggling_a1_group5.pkl": {"start": 38},
    "170407_haggling_a2_group4.pkl": {"end": 1822},
}


MANUAL_FIX = {
    "170224_haggling_a2_group7.pkl": [
        {
            "pid": 2,
            "start": 1587,
            "end": 1603,
            "jids": [0, 1, 3, 4, 5, 9, 10, 11, 15, 16, 17, 18],
        }
    ],
    "170228_haggling_a1_group5.pkl": [
        {"pid": 1, "start": 714, "end": 722, "jids": [13, 14]},
        {"pid": 1, "start": 722, "end": 750, "jids": [7, 8, 13, 14]},
    ],
    "170407_haggling_a2_group4.pkl": [
        {"pid": 1, "start": 63, "end": 69, "jids": [10, 11]},
        {"pid": 1, "start": 1531, "end": 1669, "jids": [12, 13, 14]},
    ],
}

MANUAL_FIX_LEGS = {
    "170407_haggling_a2_group4.pkl": [
        {
            "pid": 1,
            "good_frames": [
                0,
                4,
                9,
                14,
                19,
                29,
                37,
                44,
                49,
                53,
                62,
                66,
                71,
                75,
                85,
                89,
                94,
                100,
                108,
                116,
                125,
                133,
                138,
                141,
                148,
                177,
                270,
                380,
                456,
                471,
                476,
                490,
                586,
                616,
                702,
                706,
                710,
                729,
                742,
                757,
                772,
                821,
                905,
                930,
                966,
                1002,
                1051,
                1101,
                1115,
                1139,
                1152,
                1175,
                1237,
                1312,
                1403,
                1505,
                1527,
                1664,
                1693,
                1752,
                1773,
                1819,
            ],
        },
    ],
    "170224_haggling_a2_group7.pkl": [
        {
            "pid": 2,
            "good_frames": [
                0,
                5,
                10,
                15,
                18,
                20,
                24,
                26,
                30,
                32,
                34,
                36,
                39,
                41,
                46,
                53,
                58,
                61,
                64,
                67,
                71,
                74,
                76,
                79,
                81,
                87,
                91,
                95,
                102,
                114,
                121,
                173,
                212,
                247,
                251,
                256,
                261,
                273,
                279,
                283,
                290,
                296,
                306,
                444,
                488,
                602,
                758,
                884,
                1128,
                1346,
                1430,
                1435,
                1444,
                1449,
                1453,
                1460,
                1469,
                1543,
                1577,
                1613,
                1618,
                1629,
                1636,
                1641,
                1648,
                1653,
                1661,
                1668,
                1674,
                1681,
                1685,
                1688,
                1694,
                1698,
                1704,
                1710,
                1717,
                1740,
                1777,
                1819,
            ],
        }
    ],
    "170407_haggling_a1_group4.pkl": [
        {
            "pid": 2,
            "good_frames": [
                13,
                38,
                61,
                95,
                152,
                257,
                572,
                796,
                1064,
                1385,
                1810,
            ],  # noqa E501
        }
    ],
    "170407_haggling_a2_group0.pkl": [
        {
            "pid": 1,
            "good_frames": [
                0,
                22,
                31,
                33,
                38,
                46,
                49,
                90,
                99,
                109,
                136,
                194,
                241,
                349,
                386,
                488,
                519,
                549,
                657,
                681,
                743,
                785,
                815,
                865,
                935,
                1029,
                1389,
                1767,
                1785,
            ],
        }
    ],
}
social_diffusion/datasets/haggling.py
jutanke/social_diffusion
import numpy as np
import numba as nb
import numpy.linalg as la
from social_diffusion.skeleton import Skeleton
from social_diffusion.sequence import MultiPersonSequence
from social_diffusion.datasets import load_json, frames2segments
from social_diffusion import get_cache_dir, get_data_dir
from multiprocessing.pool import ThreadPool
from json import JSONDecodeError
from einops import rearrange, reduce
from os.path import isfile, isdir, join
from os import listdir

from typing import List
import logging
from time import time


logger = logging.getLogger(__name__)


PANOPTIC_SKELETON = Skeleton(
    normalizing_left_jid=6,
    normalizing_right_jid=12,
    n_joints=19,
    skeleton=[
        (0, 3),
        (3, 4),
        (4, 5),
        (0, 9),
        (9, 10),
        (10, 11),
        (0, 1),
        (0, 2),
        (2, 6),
        (2, 12),
        (6, 7),
        (7, 8),
        (12, 13),
        (13, 14),
        (15, 16),
        (15, 17),
        (17, 18),
        (1, 17),
        (1, 15),
    ],
    left_jids=[15, 16, 3, 4, 5, 6, 7, 8],
    right_jids=[17, 18, 9, 10, 11, 12, 13, 14],
    leg_joints=[7, 8, 13, 14],
)


class PanopticScene:
    """
    Contains a fixed amount of persons in the scene
    """

    def __init__(self, P: np.ndarray, name: str, skel: Skeleton):
        """
        :param P: {n_frames x n_person x jd}
        """
        self.Seq = rearrange(P, "t p j d -> t p (j d)")
        self.name = name
        self.skel = skel

    def n_person(self):
        return self.Seq.shape[1]

    def __len__(self):
        return len(self.Seq)

    def to_sequence(self) -> MultiPersonSequence:
        return MultiPersonSequence(Seq=self.Seq, skel=PANOPTIC_SKELETON)

    def trim(self, trim: int):
        """
        :param trim: cuts off at the start/end
        """
        if trim == 0:
            return self
        if trim > 0 and (2 * trim) < len(self.Seq):
            Seq = self.Seq[trim:-trim]
            name = self.name + f"_trim{trim}"
            skel = self.skel
            return PanopticScene(
                P=rearrange(Seq, "t p (j d) -> t p j d", d=3).copy(),
                name=name,
                skel=skel,
            )
        else:
            raise ValueError(
                f"Trim cannot be applied to the sequence: trim:{trim} vs Seq:{self.Seq.shape}"  # noqa E501
            )


class PanopticSequence:
    """
    Continous sequence of person(s)
    """

    def __init__(self, P: np.ndarray, M: np.ndarray, name: str):
        """
        :param P: {n_persons x n_persons x 19 x 3}
        :param M: {n_persons x n_persons}
        """
        self.name = name
        self.n_persons = P.shape[1]
        self.n_frames = P.shape[0]
        self.P = P
        self.M = M

    def __lt__(self, other):
        return len(self) < len(other)

    def __len__(self):
        return self.n_frames

    def count_persons(self):
        return self.n_persons

    def split_into_scenes_with_same_amount_of_persons(
        self,
    ) -> List[PanopticScene]:  # noqa E501
        """ """
        count = reduce(self.M, "t p -> t", "sum")
        n_max_person = self.M.shape[1]
        Seqs = []
        cur_P = []
        current_pointer = 0
        for t in range(len(count) - 1):
            if np.sum(np.abs(self.M[t] - self.M[t + 1])) > 0.1:
                cur_P = np.array(cur_P).copy()  # make sure to copy
                valid_pids = []
                for pid in range(n_max_person):
                    if self.M[t, pid] > 0.5:
                        valid_pids.append(pid)
                if len(valid_pids) == 0:
                    print("cnt", count[t], n_max_person)
                    print("self.M[t", self.M[t])
                    print("wtf?", self.name, t)
                # assert len(valid_pids) > 0
                if len(valid_pids) == 0:
                    raise ValueError("No persons..")
                valid_pids = np.array(valid_pids)
                P = np.ascontiguousarray(cur_P[:, valid_pids])
                # M = np.ones((P.shape[0], P.shape[1]), dtype=np.int64)
                Seqs.append(
                    PanopticScene(
                        P=P,
                        name=self.name + f"|esaop{current_pointer}",
                        skel=PANOPTIC_SKELETON,
                    )
                )
                current_pointer += 1
                cur_P = []
            else:
                cur_P.append(self.P[t])
        # handle last one!
        cur_P = np.array(cur_P).copy()  # make sure to copy
        valid_pids = []
        for pid in range(n_max_person):
            if self.M[t, pid] > 0.5:
                valid_pids.append(pid)
        if len(valid_pids) == 0:
            print("wtf?", self.name, t)
        if len(valid_pids) == 0:
            raise ValueError("No valid pids")

        valid_pids = np.array(valid_pids)
        P = np.ascontiguousarray(cur_P[:, valid_pids])
        Seqs.append(
            PanopticScene(
                P=P,
                name=self.name + f"|esaop{current_pointer}",
                skel=PANOPTIC_SKELETON,
            )
        )

        return Seqs


def load_poses3d(json_fname):
    """ """
    try:
        frame = int(
            json_fname.split("/")[-1].replace("body3DScene_", "").replace(".json", "")
        )
    except ValueError as e:
        logger.error(f"it seems the naming has been off: {json_fname}")
        raise e
    if not isfile(json_fname):
        raise ValueError(f"Cannot find {json_fname} :(")
    bodies = {}  # pid:
    try:
        data = load_json(json_fname)
        z_flip = np.diag([1, 1, -1])
        for body in data["bodies"]:
            pid = body["id"]
            joints19 = rearrange(
                np.array(body["joints19"], dtype=np.float32), "(j d) -> j d", j=19, d=4
            )
            joints19[:, [0, 1, 2]] = (joints19[:, [0, 2, 1]] / 100) @ z_flip
            bodies[pid] = joints19
    except JSONDecodeError:
        bodies = {}
        logger.error(f"FAILED LOADING AT FRAME {frame} ()")
        logger.error(f"FAILED JSON FNAME: {json_fname}")
    return frame, bodies


class EntirePanopticSequence:
    """
    The generic Panoptic Sequences are coarse with
    persons walking in and out of the scene
    It is even possible that the frames are not Continous!
    """

    @staticmethod
    def is_valid_scene(scene_name: str) -> bool:
        """ """
        return isdir(
            join(
                get_data_dir(), f"panoptic_studio/{scene_name}/hdPose3d_stage1_coco19"
            )  # noqa E501
        )

    def __init__(
        self,
        scene_name: str,
        local_cache=True,
    ):
        self.skel = PANOPTIC_SKELETON
        self.scene_name = scene_name
        tmp_fname = join(
            get_cache_dir(),
            f"{scene_name}.npz",
        )
        if local_cache and isfile(tmp_fname):
            o = np.load(tmp_fname)
            self.F = o["F"]
            self.P = o["P"]
            self.M = o["M"]
        else:
            logger.info(f"EntireSequence => load from Manifold: {scene_name}")
            scene_dir = join(
                get_data_dir(), f"panoptic_studio/{scene_name}/hdPose3d_stage1_coco19"
            )
            logger.info(f"Check if we need extra 'hd' path for '{scene_dir}'")
            if isdir(join(scene_dir, "hd")):
                scene_dir = join(scene_dir, "hd")
                logger.info(f"\tExtra 'hd' path exists, moving to '{scene_dir}'...")

            logger.info(f"loading and sorting files in '{scene_dir}'")
            poses3d_files = sorted(listdir(scene_dir))
            poses3d_files = [join(scene_dir, f) for f in poses3d_files]
            logger.info(f"Found #{len(poses3d_files)} files...")
            _start = time()
            with ThreadPool(20) as p:
                all_data = p.map(load_poses3d, [fname for fname in poses3d_files])

            logger.info(
                f"EntireSequence => elapsed loading from Manifold: {time() - _start}, #{ len(all_data)}"  # noqa E501
            )
            n_frames = len(all_data)
            n_max_person = 0
            for _, bodies in all_data:
                if len(bodies) > 0:
                    n_max_person = max(n_max_person, max(bodies.keys()))

            F = np.empty((n_frames,), dtype=np.int64)
            P = np.zeros((n_frames, n_max_person + 1, 19, 4), dtype=np.float32)
            M = np.zeros((n_frames, n_max_person + 1), dtype=np.int64)

            for i, (frame, bodies) in enumerate(all_data):
                F[i] = frame
                for pid, pose3d in bodies.items():
                    P[i, pid] = pose3d.reshape((19, 4))
                    M[i, pid] = 1

            self.F = F
            self.P = P
            self.M = M
            if local_cache:
                np.savez(tmp_fname, F=F, P=P, M=M)

        self.frame2index = {}
        self.index2frame = {}
        for i, frame in enumerate(self.F):
            self.frame2index[frame] = i
            self.index2frame[i] = frame

        self.interpolate_missing_limbs()

    def interpolate_missing_limbs(self):
        """ """
        n_persons = self.P.shape[1]
        with ThreadPool(n_persons) as p:
            PM_new = p.starmap(
                interpolate_missing_limbs,
                [(self.P[:, i], self.M[:, i], i) for i in range(n_persons)],
            )
            P_new = np.array([P for P, _ in PM_new], dtype=np.float32)
            M_new = np.array([M for _, M in PM_new], dtype=np.int64)
            self.P = rearrange(P_new, "p t j d -> t p j d")
            self.M = rearrange(M_new, "p t -> t p")

    def __len__(self):
        return len(self.P)

    def split_into_scenes(
        self, *, min_length=30, force_algorithm=False
    ) -> List[PanopticScene]:  # noqa E501
        scenes = []
        for sequ in self.split_into_sequences(force_algorithm=force_algorithm):
            scenes += sequ.split_into_scenes_with_same_amount_of_persons()
        return [scene for scene in scenes if len(scene) > min_length]

    def split_into_sequences(
        self, force_algorithm=False
    ) -> List[PanopticSequence]:  # noqa E501
        """ """
        sequences = []

        # check if we have manual annotations
        global SEQ_TO_ACTUAL_FRAMES
        manual_scene_names = {}
        for name in SEQ_TO_ACTUAL_FRAMES.keys():
            scene_name = name[: name.find("|")]
            if scene_name not in manual_scene_names:
                manual_scene_names[scene_name] = []
            manual_scene_names[scene_name].append(name)

        if self.scene_name in manual_scene_names and not force_algorithm:
            for name in sorted(manual_scene_names[self.scene_name]):
                start_frame, end_frame = SEQ_TO_ACTUAL_FRAMES[name]
                start = self.frame2index[start_frame]
                end = self.frame2index[end_frame]
                pids = reduce(self.M[start : end + 1], "t n -> n", "sum").nonzero()[0]
                if len(pids) == 0:  # assert len(pids) > 0
                    raise ValueError("No person...")
                P = np.ascontiguousarray(self.P[start : end + 1, pids, :, :3])
                M = np.ascontiguousarray(self.M[start : end + 1, pids])
                sequences.append(PanopticSequence(P=P, M=M, name=name))
            return sequences

        segments = frames2segments(reduce(self.M, "t n -> t", "sum").nonzero()[0])
        for idx, (start, end) in enumerate(segments):
            pids = reduce(self.M[start : end + 1], "t n -> n", "sum").nonzero()[0]
            if len(pids) == 0:  # assert len(pids) > 0
                raise ValueError("No person...")
            P = np.ascontiguousarray(self.P[start : end + 1, pids, :, :3])
            M = np.ascontiguousarray(self.M[start : end + 1, pids])
            name = self.scene_name + f"|{idx}"
            sequences.append(PanopticSequence(P=P, M=M, name=name))
        return sequences


# === PREPROCESSING ===


def interpolate_missing_poses(seq, m, segs, pid):
    for i in range(len(segs) - 1):
        left = segs[i, 1]
        right = segs[i + 1, 0]
        if right - left == 2:
            seq[left + 1] = (seq[left] + seq[right]) / 2
            m[left + 1] = 1
        else:
            print(f"\n@pid {pid}, nopy nope", segs.shape, left, right)

    return seq, m


def interpolate_missing_limbs(seq, m, pid):
    """
    :param {n_frames x 19 x 4}
    :m {n_frames}
    """
    segs = np.array(frames2segments(m.nonzero()[0]), dtype=np.int64)

    if len(segs) > 1:
        seq, m = interpolate_missing_poses(seq=seq, m=m, segs=segs, pid=pid)
        segs = np.array(frames2segments(m.nonzero()[0]), dtype=np.int64)

    if len(segs) != 1:
        raise ValueError(
            f"afaik all persons are entering only once... {segs.shape} @pid{pid}"
        )  # noqa E501

    important_jid_connections = {
        3: np.array([0, 9], dtype=np.int64),
        4: np.array([3], dtype=np.int64),
        5: np.array([4], dtype=np.int64),
        9: np.array([0, 3], dtype=np.int64),
        10: np.array([9], dtype=np.int64),
        11: np.array([10], dtype=np.int64),
        6: np.array([2, 12], dtype=np.int64),
        7: np.array([6], dtype=np.int64),
        8: np.array([7], dtype=np.int64),
        12: np.array([2, 6], dtype=np.int64),
        13: np.array([12], dtype=np.int64),
        14: np.array([13], dtype=np.int64),
        15: np.array([0, 1, 16, 17, 18], dtype=np.int64),
        16: np.array([0, 1, 15, 17, 18], dtype=np.int64),
        17: np.array([0, 1, 15, 16, 18], dtype=np.int64),
        18: np.array([0, 1, 15, 16, 17], dtype=np.int64),
    }

    for target_jid, neighbors in important_jid_connections.items():
        seq = _mark_weird_joints_as_missing(seq, segs, target_jid, neighbors)

    seq = _interpolate_missing_limbs(seq, segs)
    seq = _head_velocity_tracking(seq, segs)
    return seq, m


@nb.njit(nb.float32[:, :, :](nb.float32[:, :, :], nb.int64[:, :]), nogil=True)
def _head_velocity_tracking(seq, segments):
    head_jids = np.array([1, 18, 17, 16, 15], dtype=np.int64)
    n_segs = len(segments)
    vel = np.zeros((3,), dtype=np.float32)
    vel_hits = np.zeros((1,), dtype=np.float32)
    for i in range(n_segs):
        start = segments[i, 0]
        end = segments[i, 1]
        for t in range(start + 1, end + 1):
            vel = vel * 0  # reset velocity
            vel_hits = vel_hits * 0
            for jid in head_jids:
                if seq[t, jid, 3] > 0.08 and seq[t - 1, jid, 3] > 0.08:
                    a = seq[t, jid, :3]
                    b = seq[t - 1, jid, :3]
                    dirvec = b - a
                    vel = vel + dirvec
                    vel_hits = vel_hits + 1
            if vel_hits[0] > 0:
                vel = vel / vel_hits
            for jid in head_jids:
                if seq[t, jid, 3] < 0.1:
                    seq[t, jid, :3] = seq[t - 1, jid, :3] - vel
    return seq


@nb.njit(nb.float32[:, :, :](nb.float32[:, :, :], nb.int64[:, :]), nogil=True)
def _interpolate_missing_limbs(seq, segments):
    """
    :param seq: {n_frames x 19 x 4}
    :param segments: [(start, end)]
    """
    n_segs = len(segments)
    for i in range(n_segs):
        start = segments[i, 0]
        end = segments[i, 1]
        for jid in range(19):
            last_good_t = -1
            for t in range(start, end + 1):
                if seq[t, jid, 3] > 0:
                    last_good_t = t
                else:
                    for t_lookahead in range(t + 1, end + 1):
                        if seq[t_lookahead, jid, 3] > 0:
                            dist = float(t_lookahead - last_good_t)
                            dist_left = float(t - last_good_t)
                            dist_right = float(t_lookahead - t)

                            w_left = dist_right / dist
                            w_right = dist_left / dist

                            for d in range(3):
                                seq[t, jid, d] = (
                                    seq[last_good_t, jid, d] * w_left
                                    + seq[t_lookahead, jid, d] * w_right
                                )
                            break
    return seq


@nb.njit(
    nb.float32[:, :, :](nb.float32[:, :, :], nb.int64[:, :], nb.int64, nb.int64[:]),
    nogil=True,
)
def _mark_weird_joints_as_missing(seq, segments, target_jid, neighbor_jids):
    """
    :param seq: {n_frames x 19 x 4}
    :param segments: [(start, end)]
    """
    n_segs = len(segments)
    n_neighbors = len(neighbor_jids)

    hits = np.zeros((n_neighbors), dtype=nb.float64)
    acc_distances = np.zeros((n_neighbors), dtype=nb.float64)

    # step 0: kick out all values that are below 0.1
    for i in range(n_segs):
        start = segments[i, 0]
        end = segments[i, 1]
        for t in range(start, end + 1):
            for jid in range(19):
                if seq[t, jid, 3] < 0.08:
                    seq[t, jid, 3] = -1

    # step 1: calculate average dist
    for i in range(n_segs):
        start = segments[i, 0]
        end = segments[i, 1]
        for t in range(start, end + 1):
            if seq[t, target_jid, 3] > 0:
                pt1 = seq[t, target_jid, :3]
                for n, neighbor_jid in enumerate(neighbor_jids):
                    if seq[t, neighbor_jid, 3] > 0:
                        pt2 = seq[t, neighbor_jid, :3]
                        d = la.norm(pt2 - pt1)
                        acc_distances[n] += d
                        hits[n] += 1

        # step 2: mark as "weird" if neighbor distance too large
        leeway = 1.3
        for n in range(n_neighbors):
            if hits[n] > 0:
                acc_distances[n] = (acc_distances[n] / hits[n]) * leeway

        weird_factor_step = 1.0 / n_neighbors
        for i in range(n_segs):
            start = segments[i, 0]
            end = segments[i, 1]
            for t in range(start, end + 1):
                if seq[t, target_jid, 3] > 0:
                    pt1 = seq[t, target_jid, :3]

                    weird_factor = 0.0  # not weird <-- 0....0.5....1 --> weird
                    for n, neighbor_jid in enumerate(neighbor_jids):
                        if hits[n] > 0:
                            if seq[t, neighbor_jid, 3] > 0:
                                if seq[t, neighbor_jid, 3] > 0:
                                    pt2 = seq[t, neighbor_jid, :3]
                                    d = la.norm(pt2 - pt1)
                                    if d > acc_distances[n]:
                                        weird_factor += weird_factor_step
                                else:
                                    weird_factor += weird_factor_step * 0.5
                    if weird_factor > 0.5:
                        seq[t, target_jid, 3] = -1
    return seq


# some of the sequences are bady cut: we fix that here
SEQ_TO_ACTUAL_FRAMES = {
    "160224_haggling1|0": (224, 2129),
    "160224_haggling1|1": (2487, 4342),
    "160224_haggling1|2": (4706, 6831),
    "160224_haggling1|3": (7109, 8761),
    "160226_haggling1|0": (130, 1973),
    "160226_haggling1|1": (2224, 4035),
    "160226_haggling1|2": (4226, 5715),
    "160226_haggling1|3": (5905, 7902),
    "160226_haggling1|4": (8187, 10009),
    "160226_haggling1|5": (10206, 11593),
    "160422_haggling1|0": (333, 2501),
    "160422_haggling1|1": (2811, 4722),
    "160422_haggling1|2": (4958, 7024),
    "160422_haggling1|3": (7188, 8943),
    "160422_haggling1|4": (9149, 10936),
    "160422_haggling1|5": (11290, 13185),
    "160906_band1|0": (185, 1939),
    "160906_band3|0": (162, 7500),
    "161029_build1|0": (191, 9005),
    "161029_piano1|0": (1787, 7531),  # cut out the child
    "161029_sports1|0": (258, 730),
    "161029_sports1|1": (778, 2677),
    "161029_sports1|2": (2755, 4492),
    "161029_sports1|3": (4905, 5371),
    "161029_sports1|4": (5590, 6045),
    "161029_sports1|5": (6605, 7219),
    "161029_sports1|6": (7904, 8077),
    "161029_sports1|7": (8205, 8976),
    "161029_tools1|0": (144, 5523),
    "170221_haggling_b1|0": (1097, 3244),
    "170221_haggling_b1|1": (3518, 5409),
    "170221_haggling_b1|2": (5741, 7885),
    "170221_haggling_b1|3": (8221, 10333),
    "170221_haggling_b1|4": (10572, 12655),
    "170221_haggling_b1|5": (12929, 15055),
    "170221_haggling_b3|0": (255, 2425),
    "170221_haggling_b3|1": (2576, 4647),
    "170221_haggling_b3|2": (4915, 7023),
    "170221_haggling_b3|3": (8218, 8506),  # POSES!
    "170221_haggling_b3|4": (8909, 9521),  # POSES!
    "170221_haggling_b3|5": (9699, 10263),  # POSES!
    "170221_haggling_b3|6": (10464, 10961),  # POSES!
    "170221_haggling_b3|6": (11161, 11732),  # POSES!
    "170221_haggling_b3|7": (11936, 12466),  # POSES!
    "170221_haggling_b3|8": (12590, 13151),  # POSES!
    "170221_haggling_b3|9": (13388, 13879),  # POSES!
    "170221_haggling_b3|10": (14058, 14605),  # POSES!
    "170221_haggling_b3|11": (14681, 15115),  # POSES!
}
social_diffusion/datasets/panoptic.py
jutanke/social_diffusion
"""
https://github.com/lucidrains/denoising-diffusion-pytorch
"""
import math
from collections import namedtuple
from functools import partial

import numpy as np
import torch
import torch.nn.functional as F
from einops import reduce, repeat, rearrange
from social_diffusion.skeleton import Skeleton
from torch import nn
import numpy.linalg as la
from tqdm import tqdm
from social_diffusion.transforms.rotation import rot3d

ModelPrediction = namedtuple("ModelPrediction", ["pred_noise", "pred_x_start"])


def overwrite(inpainting_mask, P_source, P_target):
    """
    P_source -> P_target

    :param inpainting_mask: {b x t x p x jd}
    :param P_source: {b x t x p x jd}
    :param P_target: {b x t x p x jd}
    """
    assert P_source.shape == P_target.shape
    # assert P_source.shape == inpainting_mask.shape
    b, t, p, jd = P_source.shape

    flat_mask = rearrange(inpainting_mask, "b t p jd -> (b t p jd)")

    indices = np.nonzero(flat_mask)[0]

    P_source_flat = rearrange(P_source, "b t p jd -> (b t p jd)")
    P_target_flat = rearrange(P_target, "b t p jd -> (b t p jd)")

    P_target_flat[indices] = P_source_flat[indices]
    return rearrange(P_target_flat, "(b t p jd) -> b t p jd", b=b, t=t, p=p, jd=jd)


class DiffusionModel(nn.Module):
    """
    Template model for the diffusion process
    """

    def __init__(
        self,
        skel: Skeleton,
    ):
        super().__init__()
        self._skel = skel

    def skel(self):
        return self._skel

    def jd(self):
        return self._skel.n_joints * 3

    def forward(self, X_0, X_t, batch):
        raise NotImplementedError("Nope")


def exists(x):
    return x is not None


def identity(t, *args, **kwargs):
    return t


def default(val, d):
    if exists(val):
        return val
    return d() if callable(d) else d


def has_int_squareroot(num):
    return (math.sqrt(num) ** 2) == num


def extract(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def linear_beta_schedule(timesteps):
    scale = 1000 / timesteps
    beta_start = scale * 0.0001
    beta_end = scale * 0.02
    return torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)


def l2norm(t):
    return F.normalize(t, dim=-1)


def cosine_beta_schedule(timesteps, s=0.008):
    """
    cosine schedule
    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
    """
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps, dtype=torch.float64)
    alphas_cumprod = (
        torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
    )  # noqa E501
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0, 0.999)


class GaussianDiffusion(nn.Module):
    def __init__(
        self,
        model: DiffusionModel,
        *,
        n_in: int,
        n_frames: int,
        n_person: int,
        timesteps=1000,
        sampling_timesteps=None,
        loss_type="l1",
        objective="pred_noise",
        beta_schedule="cosine",
        # p2 loss weight, from https://arxiv.org/abs/2204.00227 -
        #    0 is equivalent to weight of 1 across time - 1. is recommended
        p2_loss_weight_gamma=0.0,
        p2_loss_weight_k=1,
        ddim_sampling_eta=1.0,
        # same_noise_for_all_dims=False,
        # additional_loss_fn=None,
    ):
        super().__init__()
        self.model = model

        # self.same_noise_for_all_dims = same_noise_for_all_dims

        # if additional_loss_fn is None:
        #     self.additional_loss_fn = lambda gt, pred: 0
        # else:
        #     self.additional_loss_fn = additional_loss_fn

        self.n_in = n_in
        assert n_frames > n_in

        self.n_frames = n_frames
        self.n_person = n_person

        self.objective = objective

        assert objective in {
            "pred_noise",
            "pred_x0",
        }, "objective must be either pred_noise (predict noise) or pred_x0(predict image start)"  # noqa E501

        if beta_schedule == "linear":
            betas = linear_beta_schedule(timesteps)
        elif beta_schedule == "cosine":
            betas = cosine_beta_schedule(timesteps)
        else:
            raise ValueError(f"unknown beta schedule {beta_schedule}")

        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, axis=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)

        (timesteps,) = betas.shape
        self.num_timesteps = int(timesteps)
        self.loss_type = loss_type

        # sampling related parameters

        self.sampling_timesteps = default(
            sampling_timesteps, timesteps
        )  # default num sampling timesteps to number of timesteps at training

        assert self.sampling_timesteps <= timesteps
        self.is_ddim_sampling = self.sampling_timesteps < timesteps
        self.ddim_sampling_eta = ddim_sampling_eta

        # helper function to register buffer from float64 to float32

        def register_buffer(name, val):
            return self.register_buffer(name, val.to(torch.float32))

        register_buffer("betas", betas)
        register_buffer("alphas_cumprod", alphas_cumprod)
        register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)

        # calculations for diffusion q(x_t | x_{t-1}) and others

        register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        register_buffer(
            "sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod)
        )
        register_buffer(
            "log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod)
        )  # noqa E501
        register_buffer(
            "sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod)
        )  # noqa E501
        register_buffer(
            "sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1)
        )

        # calculations for posterior q(x_{t-1} | x_t, x_0)

        posterior_variance = (
            betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        )

        # noqa E501, above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)

        register_buffer("posterior_variance", posterior_variance)

        # noqa E501, below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain

        register_buffer(
            "posterior_log_variance_clipped",
            torch.log(posterior_variance.clamp(min=1e-20)),
        )
        register_buffer(
            "posterior_mean_coef1",
            betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
        )
        register_buffer(
            "posterior_mean_coef2",
            (1.0 - alphas_cumprod_prev)
            * torch.sqrt(alphas)
            / (1.0 - alphas_cumprod),  # noqa E501
        )

        # calculate p2 reweighting

        register_buffer(
            "p2_loss_weight",
            (p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod))
            ** -p2_loss_weight_gamma,
        )

    def predict_start_from_noise(self, x_t, t, noise):
        return (
            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
        )

    def predict_noise_from_start(self, x_t, t, x0):
        return (
            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0
        ) / extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)

    def q_posterior(self, x_start, x_t, t):
        posterior_mean = (
            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
            + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = extract(
            self.posterior_log_variance_clipped, t, x_t.shape
        )
        return (
            posterior_mean,
            posterior_variance,
            posterior_log_variance_clipped,
        )  # noqa E501

    def model_predictions(self, x, t, x_0, batch, clip_x_start=False):
        model_output = self.model(X_0=x_0, X_t=x, time=t, batch=batch)
        maybe_clip = (
            partial(torch.clamp, min=-1.0, max=1.0)
            if clip_x_start
            else identity  # noqa E501
        )

        if self.objective == "pred_noise":
            pred_noise = model_output
            x_start = self.predict_start_from_noise(x, t, pred_noise)
            x_start = maybe_clip(x_start)

        elif self.objective == "pred_x0":
            x_start = model_output
            x_start = maybe_clip(x_start)
            pred_noise = self.predict_noise_from_start(x, t, x_start)

        return ModelPrediction(pred_noise, x_start)

    # def p_mean_variance(self, x, t, x_0, batch, clip_denoised=True):
    def p_mean_variance(self, x, t, x_0, batch, clip_denoised=False):
        preds = self.model_predictions(x=x, t=t, x_0=x_0, batch=batch)
        x_start = preds.pred_x_start

        if clip_denoised:
            x_start.clamp_(-1.0, 1.0)

        (
            model_mean,
            posterior_variance,
            posterior_log_variance,
        ) = self.q_posterior(  # noqa E501
            x_start=x_start, x_t=x, t=t
        )
        return model_mean, posterior_variance, posterior_log_variance, x_start

    @torch.no_grad()
    def p_sample(
        self,
        x,
        t: int,
        x_0: torch.Tensor,
        batch: torch.Tensor,
        clip_denoised=False,
        init_rand=None,
    ):
        batched_times = torch.full(
            (x.shape[0],), t, device=x.device, dtype=torch.long
        )  # noqa E501

        model_mean, _, model_log_variance, x_start = self.p_mean_variance(
            x=x,
            t=batched_times,
            x_0=x_0,
            clip_denoised=clip_denoised,
            batch=batch,
        )

        if init_rand is None:
            noise = (
                torch.randn_like(x) if t > 0 else 0.0
            )  # no noise if t == 0   # noqa E501
        else:
            noise = init_rand
        pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
        return pred_img, x_start

    @torch.no_grad()
    def p_sample_loop_with_input(
        self,
        batch,
        return_all_steps=False,
        # add_gt=True,
        add_gt=False,
        linear_blending=None,
        init_rand=None,
        return_noise=False,
        n_diffusion_steps=-1,
        ds_train=None,
    ):
        """
        :param batch: {
            "X_0": ...
        }
        """
        device = self.betas.device

        P_out = batch["X_0"]
        P_out = P_out.to(device)

        Imgs = []

        if init_rand is None:
            img = torch.randn(P_out.shape, device=device)
        else:
            img = init_rand

        init_noise = img.clone()

        img = img.contiguous()

        # img_last = None

        img[:, : self.n_in] = P_out[:, : self.n_in]
        if linear_blending is not None:
            lin_blend_start, lin_blend_end = linear_blending
            img[:, lin_blend_start:lin_blend_end] = P_out[
                :, lin_blend_start:lin_blend_end
            ]

        if return_all_steps:
            Imgs.append(img.clone().cpu().numpy())

        x_start = None

        def sqish(t, num_timesteps, shift=0):
            x = num_timesteps / 2 - t + shift
            return 1 - (1 / (1 + np.exp(-x * 0.01)))

        n_diffusion_steps = (
            n_diffusion_steps if n_diffusion_steps > 0 else self.num_timesteps
        )
        for t in tqdm(
            reversed(range(0, n_diffusion_steps)),
            desc="sampling loop time step",
            total=n_diffusion_steps,
            position=0,
            leave=True,
        ):
            img, x_start = self.p_sample(
                x=img, t=t, x_0=P_out, init_rand=init_rand, batch=batch
            )
            # img_last = img.copy()

            # adjust orientation!
            if ds_train is not None:
                n_person = P_out.size(2)
                n_frames = P_out.size(1)

                img_ = []
                n_samples = img.shape[0]
                for ii in range(n_samples):
                    img_seq = ds_train.undo_standardization(img[ii])

                    p_out_seq = ds_train.undo_standardization(
                        P_out[ii].detach().cpu().numpy()
                    )

                    xy_pred = img_seq[self.n_in - 1, :, 3:5]
                    xy_pred = xy_pred / la.norm(xy_pred, axis=1, keepdims=True)

                    x_pred = xy_pred[:, 0]
                    y_pred = xy_pred[:, 1]
                    rads_pred = np.arctan2(y_pred, x_pred)

                    x_real = p_out_seq[self.n_in - 1, :, 3]
                    y_real = p_out_seq[self.n_in - 1, :, 4]

                    rads_real = np.arctan2(y_real, x_real)

                    rads_pred2real = rads_pred - rads_real

                    for iii in range(n_person):
                        translation = (
                            p_out_seq[self.n_in - 1, iii, :3]
                            - img_seq[self.n_in - 1, iii, :3]
                        )
                        rad = rads_pred2real[iii]
                        if rad < 0:
                            rad += 2 * np.pi
                        R = np.transpose(rot3d(0, 0, rad))

                        for frame in range(n_frames):
                            pred = np.zeros((3,), dtype=np.float32)
                            pred[:2] = img_seq[frame, iii, 3:5]
                            img_seq[frame, iii, 3:6] = R @ pred
                            img_seq[frame, iii, :3] += translation

                    img_seq = ds_train.do_standardization(img_seq)
                    img_.append(img_seq)

                img = torch.from_numpy(np.array(img_)).to(P_out.device)

            if return_all_steps:
                Imgs.append(img.clone().cpu().numpy())
            img[:, : self.n_in] = P_out[:, : self.n_in]

            if linear_blending is not None:
                lin_blend_start, lin_blend_end = linear_blending
                blend_value = np.linspace(
                    start=0.95,
                    stop=0.0001,
                    num=lin_blend_end - lin_blend_start,  # noqa E501
                )
                for i, blend_t in enumerate(
                    range(lin_blend_start, lin_blend_end)
                ):  # noqa E501
                    pred = P_out[:, blend_t]
                    diff = img[:, blend_t]

                    pred_w = blend_value[i]
                    diff_w = 1.0 - pred_w

                    img[:, blend_t] = pred * pred_w + diff * diff_w

            # img[:, lin_blend_start:lin_blend_end] = P_out[
            #     :, lin_blend_start:lin_blend_end
            # ]

        if add_gt:
            img[:, : self.n_in] = P_out[:, : self.n_in]
            if linear_blending is not None:
                lin_blend_start, lin_blend_end = linear_blending
                blend_value = np.linspace(
                    start=0.95,
                    stop=0.0001,
                    num=lin_blend_end - lin_blend_start,  # noqa E501
                )
                for i, blend_t in enumerate(
                    range(lin_blend_start, lin_blend_end)
                ):  # noqa E501
                    pred = P_out[:, blend_t]
                    diff = img[:, blend_t]

                    pred_w = blend_value[i]
                    diff_w = 1.0 - pred_w

                    img[:, blend_t] = pred * pred_w + diff * diff_w
        if return_all_steps:
            assert not return_noise
            Imgs.append(img.cpu().numpy())
            return Imgs

        if return_noise:
            return img, init_noise

        return img

    @torch.no_grad()
    def p_sample_loop(self, shape):
        _, device = shape[0], self.betas.device

        img = torch.randn(shape, device=device)

        x_start = None

        for t in tqdm(
            reversed(range(0, self.num_timesteps)),
            desc="sampling loop time step",
            total=self.num_timesteps,
        ):
            img, x_start = self.p_sample(img, t, x_0=img, batch={})

        return img

    @torch.no_grad()
    def inpainting_sample_with_input(
        self, batch, *, n_diffusion_steps=1000, ds_train=None
    ):
        return self.p_inpatining_sample_loop_with_input(
            batch=batch, n_diffusion_steps=n_diffusion_steps, ds_train=ds_train
        )

    @torch.no_grad()
    def p_inpatining_sample_loop_with_input(
        self, batch, *, n_diffusion_steps=1000, ds_train=None
    ):
        device = self.betas.device
        P_out = batch["X_0"]
        inpainting_mask = batch["inpainting_mask"]
        P_out = P_out.to(device)
        img = torch.randn(P_out.shape, device=device)
        # def overwrite(inpainting_mask, P_source, P_target):
        img = overwrite(
            inpainting_mask=inpainting_mask, P_source=P_out, P_target=img
        ).contiguous()

        n_diffusion_steps = (
            n_diffusion_steps if n_diffusion_steps > 0 else self.num_timesteps
        )
        for t in tqdm(
            reversed(range(0, n_diffusion_steps)),
            desc="sampling loop time step",
            total=n_diffusion_steps,
            position=0,
            leave=True,
        ):
            img, x_start = self.p_sample(
                x=img, t=t, x_0=P_out, init_rand=None, batch=batch
            )
            img = overwrite(
                inpainting_mask=inpainting_mask, P_source=P_out, P_target=img
            ).contiguous()

        return img.detach().cpu().numpy()

    @torch.no_grad()
    def sample_with_input(
        self,
        batch,
        linear_blending=None,
        return_all_steps=False,
        init_rand=None,
        return_noise=False,
        n_diffusion_steps=-1,
        ds_train=None,
    ):
        return self.p_sample_loop_with_input(
            batch,
            linear_blending=linear_blending,
            return_all_steps=return_all_steps,
            init_rand=init_rand,
            return_noise=return_noise,
            n_diffusion_steps=n_diffusion_steps,  # JULIAN
            add_gt=False,
            ds_train=ds_train,
        )

    @torch.no_grad()
    def sample(self, batch_size=8, n_frames=None):
        if n_frames is None:
            n_frames = self.n_frames
        sample_fn = (
            self.p_sample_loop
            if not self.is_ddim_sampling
            else self.ddim_sample  # noqa E501
        )
        return sample_fn(
            (batch_size, n_frames, self.n_person, self.model.jd())
        )  # noqa E501

    def q_sample(self, x_start, t, noise=None):
        """
        :param x_start: {n_batch x n_frames x n_person x dim}
        """
        noise = default(noise, lambda: torch.randn_like(x_start))

        return (
            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * noise  # noqa E501
        )

    @property
    def loss_fn(self):
        if self.loss_type == "l1":
            return F.l1_loss
        elif self.loss_type == "l2":
            return F.mse_loss
        else:
            raise ValueError(f"invalid loss type {self.loss_type}")

    def p_losses(self, P_out, t, batch, noise=None):
        """
        :param P_out: {n_frames x n_frames x 3 x 57}
        """
        x_start = P_out

        noise = default(noise, lambda: torch.randn_like(x_start))

        x = self.q_sample(x_start=x_start, t=t, noise=noise)

        model_out = self.model(X_0=x_start, X_t=x, time=t, batch=batch)

        if self.objective == "pred_noise":
            target = noise
        elif self.objective == "pred_x0":
            target = x_start
        else:
            raise ValueError(f"unknown objective {self.objective}")

        loss = self.loss_fn(model_out, target, reduction="none")
        # loss = reduce(loss, "b d t j -> t", "mean")

        weights_per_j = torch.from_numpy(
            self.model.skel().get_loss_weight()
        ).to(  # noqa E501
            loss.device
        )

        loss = reduce(loss, "b d t j -> j", "mean")
        loss = loss * weights_per_j
        # b t p jd

        loss = loss.mean()

        final_loss = loss
        return final_loss, model_out

    def forward(self, batch, *args, **kwargs):
        """
        :param P_in: {n_batch x n_frames x 3 x 57}
        """
        P_out = batch["X_0"]
        b = P_out.size(0)
        device = P_out.device

        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()

        return self.p_losses(P_out, t, batch, *args, **kwargs)

    def inpainting(
        self,
        P,
        inpainting_mask,
        ds_train,
        *,
        device=None,
        n_samples=1,
    ):
        """
        :param P:
        """

        skel = self.model.skel()
        if device is None:
            device = torch.device("cuda:0")
        assert isinstance(P, np.ndarray)
        if len(P.shape) != 3 or P.shape[2] != skel.jd_without_normalization():
            raise ValueError(f"Weird shape: {P.shape}")

        n_frames = P.shape[0]
        # if n_in == -1:
        #     n_in = self.n_in
        n_person = P.shape[1]

        P_norm = np.zeros((n_frames, n_person, skel.jd()), dtype=np.float32)
        P_norm, undo_data = ds_train.create_batch_entry(P)

        # if after_create_batch_entry_fn is not None:
        #     P_norm = after_create_batch_entry_fn(P_norm)

        P_norm = ds_train.do_standardization(P_norm)

        # P_norm[n_in:] = 0
        # for t in range(n_in, len(P_norm)):
        #     P_norm[t] = P_norm[n_in - 1]

        P_norm = torch.from_numpy(P_norm).to(device).unsqueeze(0)

        if n_samples > 1:
            P_norm = repeat(P_norm, "b t p jd -> (s b) t p jd", s=n_samples)

        inpainting_mask_full = np.ones(P_norm.shape)
        inpainting_mask_full[:, :, :, 6:] = rearrange(
            inpainting_mask, "t p jd -> 1 t p jd"
        )
        # inpainting_mask_full = torch.from_numpy(inpainting_mask_full).float().to(device)

        O = self.inpainting_sample_with_input(
            batch={"X_0": P_norm, "inpainting_mask": inpainting_mask_full},
            # return_noise=False,
            n_diffusion_steps=1000,
            # return_all_steps=False,
            ds_train=ds_train,
        )

        print("P_norm", P_norm.shape)
        print("O", O.shape)
        for ii in range(len(O)):
            O[ii] = ds_train.undo_standardization(O[ii])

        P_out = []
        for ii in range(len(O)):
            P_out.append(ds_train.undo_batch_entry(O[ii], undo_data))

        P_out = np.array(P_out, dtype=np.float32)

        return P_out

    def predict(
        self,
        P,
        ds_train,
        *,
        device=None,
        n_in=-1,
        n_samples=1,
        after_create_batch_entry_fn=None,
        before_undo_standardization_fn=None,
        before_undo_batch_entry_fn=None,
        return_noise=False,
        n_diffusion_steps=-1,
    ):
        """
        actually predicts a pose!
        :param after_create_batch_entry_fn:
            def after_create_batch_entry_fn(P) -> P
        :param before_undo_standardization_fn:
            def before_undo_standardization_fn(P) -> P
        :param before_undo_batch_entry_fn:
            def before_undo_batch_entry_fn(P) -> P
        """
        skel = self.model.skel()
        if device is None:
            device = torch.device("cuda:0")
        assert isinstance(P, np.ndarray)
        if len(P.shape) != 3 or P.shape[2] != skel.jd_without_normalization():
            raise ValueError(f"Weird shape: {P.shape}")

        n_frames = P.shape[0]
        if n_in == -1:
            n_in = self.n_in
        n_person = P.shape[1]

        P_norm = np.zeros((n_frames, n_person, skel.jd()), dtype=np.float32)
        P_norm[:n_in], undo_data = ds_train.create_batch_entry(P[:n_in])

        if after_create_batch_entry_fn is not None:
            P_norm = after_create_batch_entry_fn(P_norm)

        P_norm = ds_train.do_standardization(P_norm)

        P_norm[n_in:] = 0
        for t in range(n_in, len(P_norm)):
            P_norm[t] = P_norm[n_in - 1]

        P_norm = torch.from_numpy(P_norm).to(device).unsqueeze(0)

        if n_samples > 1:
            P_norm = repeat(P_norm, "b t p jd -> (s b) t p jd", s=n_samples)

        Os = self.sample_with_input(
            batch={"X_0": P_norm},
            return_noise=False,
            n_diffusion_steps=n_diffusion_steps,
            return_all_steps=True,
            ds_train=ds_train,
        )
        O = Os[-1]  # noqa E741

        if before_undo_standardization_fn is not None:
            for ii in range(len(O)):
                O[ii] = before_undo_standardization_fn(O[ii])

        for ii in range(len(O)):
            O[ii] = ds_train.undo_standardization(O[ii])

        if before_undo_batch_entry_fn is not None:
            for ii in range(len(O)):
                O[ii] = before_undo_batch_entry_fn(O[ii])

        P_out = []
        for ii in range(len(O)):
            P_out.append(ds_train.undo_batch_entry(O[ii], undo_data))

        P_out = np.array(P_out, dtype=np.float32)

        return P_out
social_diffusion/diffusion.py
jutanke/social_diffusion
"""
https://github.com/lucidrains/ema-pytorch
"""


import copy

import torch
from torch import nn


def exists(val):
    return val is not None


def clamp(value, min_value=None, max_value=None):
    assert exists(min_value) or exists(max_value)
    if exists(min_value):
        value = max(value, min_value)

    if exists(max_value):
        value = min(value, max_value)

    return value


class EMA(nn.Module):
    """
    Implements exponential moving average shadowing for your model.
    Utilizes an inverse decay schedule to manage longer term training runs.
    By adjusting the power, you can control how fast EMA will ramp up to your
    specified beta.
    @crowsonkb's notes on EMA Warmup:
    If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are
    good values for models you plan to train for a million or more steps
    (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps), gamma=1,
    power=3/4 for models you plan to train for less (reaches decay factor
    0.999 at 10K steps, 0.9999 at 215.4k steps).
    Args:
        inv_gamma (float): Inverse multiplicative factor of EMA warmup.
        Default: 1.
        power (float): Exponential factor of EMA warmup. Default: 1.
        min_value (float): The minimum EMA decay rate. Default: 0.
    """

    def __init__(
        self,
        model,
        ema_model=None,  # noqa E501 if your model has lazylinears or other types of non-deepcopyable modules, you can pass in your own ema model
        beta=0.9999,
        update_after_step=100,
        update_every=10,
        inv_gamma=1.0,
        power=2 / 3,
        min_value=0.0,
        param_or_buffer_names_no_ema=set(),
        ignore_names=set(),
    ):
        super().__init__()
        self.beta = beta
        self.online_model = model

        self.ema_model = ema_model

        if not exists(self.ema_model):
            try:
                self.ema_model = copy.deepcopy(model)
            except:  # noqa E722
                print(
                    "Your model was not copyable. Please make sure you are not using any LazyLinear"  # noqa E501
                )
                exit()

        self.ema_model.requires_grad_(False)

        self.parameter_names = {
            name
            for name, param in self.ema_model.named_parameters()
            if param.dtype == torch.float
        }
        self.buffer_names = {
            name
            for name, buffer in self.ema_model.named_buffers()
            if buffer.dtype == torch.float
        }

        self.update_every = update_every
        self.update_after_step = update_after_step

        self.inv_gamma = inv_gamma
        self.power = power
        self.min_value = min_value

        assert isinstance(param_or_buffer_names_no_ema, (set, list))
        self.param_or_buffer_names_no_ema = (
            param_or_buffer_names_no_ema  # parameter or buffer
        )

        self.ignore_names = ignore_names

        self.register_buffer("initted", torch.Tensor([False]))
        self.register_buffer("step", torch.tensor([0]))

    def restore_ema_model_device(self):
        device = self.initted.device
        self.ema_model.to(device)

    def get_params_iter(self, model):
        for name, param in model.named_parameters():
            if name not in self.parameter_names:
                continue
            yield name, param

    def get_buffers_iter(self, model):
        for name, buffer in model.named_parameters():
            if name not in self.buffer_names:
                continue
            yield name, buffer

    def copy_params_from_model_to_ema(self):
        for (_, ma_params), (_, current_params) in zip(
            self.get_params_iter(self.ema_model),
            self.get_params_iter(self.online_model),
        ):
            ma_params.data.copy_(current_params.data)

        for (_, ma_buffers), (_, current_buffers) in zip(
            self.get_buffers_iter(self.ema_model),
            self.get_buffers_iter(self.online_model),
        ):
            ma_buffers.data.copy_(current_buffers.data)

    def get_current_decay(self):
        epoch = clamp(
            self.step.item() - self.update_after_step - 1, min_value=0.0
        )  # noqa E501
        value = 1 - (1 + epoch / self.inv_gamma) ** -self.power

        if epoch <= 0:
            return 0.0

        return clamp(value, min_value=self.min_value, max_value=self.beta)

    def update(self):
        step = self.step.item()
        self.step += 1

        if (step % self.update_every) != 0:
            return

        if step <= self.update_after_step:
            self.copy_params_from_model_to_ema()
            return

        if not self.initted.item():
            self.copy_params_from_model_to_ema()
            self.initted.data.copy_(torch.Tensor([True]))

        self.update_moving_average(self.ema_model, self.online_model)

    @torch.no_grad()
    def update_moving_average(self, ma_model, current_model):
        current_decay = self.get_current_decay()

        for (name, current_params), (_, ma_params) in zip(
            self.get_params_iter(current_model), self.get_params_iter(ma_model)
        ):
            if name in self.ignore_names:
                continue

            if name in self.param_or_buffer_names_no_ema:
                ma_params.data.copy_(current_params.data)
                continue

            difference = ma_params.data - current_params.data
            difference.mul_(1.0 - current_decay)
            ma_params.sub_(difference)

        for (name, current_buffer), (_, ma_buffer) in zip(
            self.get_buffers_iter(current_model),
            self.get_buffers_iter(ma_model),  # noqa E501
        ):
            if name in self.ignore_names:
                continue

            if name in self.param_or_buffer_names_no_ema:
                ma_buffer.data.copy_(current_buffer.data)
                continue

            difference = ma_buffer - current_buffer
            difference.mul_(1.0 - current_decay)
            ma_buffer.sub_(difference)

    def __call__(self, *args, **kwargs):
        return self.ema_model(*args, **kwargs)
social_diffusion/ema.py
jutanke/social_diffusion
from os.path import join
from social_diffusion.datasets.haggling import get_haggling_local_dataset
from social_diffusion.train import load_socdiff
import numpy as np
from einops import repeat
from social_diffusion import get_output_dir
from social_diffusion.sequence import Sequence
import torch


def predict(ema_diffusion, ds_train, Seq: Sequence, *, n_samples=8):
    INPUT_FRAME_10PERC = 178

    P = np.zeros((2048, 3, 57), dtype=np.float32)
    P[:128] = Seq[INPUT_FRAME_10PERC - 128 : INPUT_FRAME_10PERC]  # noqa E203
    Seq_hat = ema_diffusion.ema_model.predict(
        P, ds_train=ds_train, n_samples=n_samples
    )  # noqa E501
    Seq_hat_init = repeat(
        Seq[: INPUT_FRAME_10PERC - 128], "t p jd -> s t p jd", s=n_samples
    )  # noqa E501
    return np.concatenate([Seq_hat_init, Seq_hat], axis=1)


def load_trained_model():
    n_in = 128
    n_frames = 512
    n_person = 3
    hidden_dim = 128 + 64

    model_dir = join(
        get_output_dir(),
        f"models/socdiff_haggling_nin{n_in}_nframes{n_frames}_hidden{hidden_dim}",  # noqa E501
    )
    device = torch.device("cuda")

    ds_train, _, skel = get_haggling_local_dataset(
        n_frames=n_frames, n_in=n_in
    )  # noqa E501

    _, ema_diffusion = load_socdiff(
        model_dir=model_dir,
        n_in=n_in,
        n_frames=n_frames,
        hidden_dim=hidden_dim,
        epoch=50,
        skel=skel,
        n_person=n_person,
        device=device,
    )

    return ds_train, ema_diffusion

social_diffusion/eval_utils.py
jutanke/social_diffusion
from os.path import join
import math
import os

import numpy as np
import pytorch_lightning as pl
import torch
import torch.nn.functional as F
import torch.nn as nn
from einops import reduce, rearrange
from social_diffusion.transforms.transforms import (
    get_normalization,
    apply_normalization_to_seq,
)


class WhoIsBuyerDecider:
    def __init__(self, n_frames=512, device="cuda"):
        """ """
        self.n_frames = n_frames
        fname_local = join(os.getcwd(), "models/whoisbuyer.ckpt")

        self.model = WhoIsTheBuyerModule.load_from_checkpoint(fname_local)
        self.model.eval().to(device)
        self.device = device

    def who_is_the_buyer(self, P, ss=50):
        """
        :param P: {n_frames x 3 x 57}
        """
        selected = []
        n_frames = len(P)
        for t in range(0, n_frames - self.n_frames, ss):
            P_chunk = P[t : t + self.n_frames]  # noqa E203
            pid = self.eval_chunk(P_chunk)
            selected.append(pid)

        U, C = np.unique(selected, return_counts=True)
        best_u = -1
        best_c = -1
        for u, c in zip(U, C):
            if c > best_c:
                best_c = c
                best_u = u
        return best_u

    def eval_chunk(self, P):
        """
        :param P: {n_frames x 3 x 57}
        """
        assert (
            len(P.shape) == 3
            and P.shape[0] == self.n_frames
            and P.shape[1] == 3
            and P.shape[2] == 57
        ), f"weird P shape: {P.shape}"

        index_rearrange, B = ensure_clockwise_order(P)

        B = torch.from_numpy(B).to(self.device)

        with torch.no_grad():
            pred = self.model.model(B).cpu().numpy()

        return index_rearrange[np.argmax(pred)]


def ensure_clockwise_order(P):
    """
    :param P: {n_frames x 3 x 57}
    :returns:
        {3 x n_frames x 3 x 57}
    """
    assert (
        len(P.shape) == 3 and P.shape[1] == 3 and P.shape[2] == 57
    ), f"weird P shape: {P.shape}"
    P = rearrange(P, "t p (j d) -> t p j d", j=19, d=3)
    A = P[:, 0]
    B = P[:, 1]
    C = P[:, 2]

    # we don't know the order (who's left/right of whom)
    # if we follow clock-wise we want the following order:
    # A --> B --> C --> A
    # step 1: figure out order and make sure its the above order!
    mean_hip = reduce(A[:, [6, 12]], "t j d -> j d", "mean")
    mu, R = get_normalization(mean_hip[0], mean_hip[1])

    B_in_A = apply_normalization_to_seq(B, mu=mu, R=R)
    C_in_A = apply_normalization_to_seq(C, mu=mu, R=R)

    # # the person with the lower x-value is the left-hand person
    mean_B = reduce(B_in_A[:, [6, 12]], "t j d -> d", "mean")
    mean_C = reduce(C_in_A[:, [6, 12]], "t j d -> d", "mean")

    index_rearrange = [0, 1, 2]

    if mean_B[0] > mean_C[0]:
        B, C = C, B
        index_rearrange[1] = 2
        index_rearrange[2] = 1

    SeqABC = normalize_data(target_seq=A, left_seq=B, right_seq=C)
    SeqBCA = normalize_data(target_seq=B, left_seq=C, right_seq=A)
    SeqCAB = normalize_data(target_seq=C, left_seq=A, right_seq=B)

    SeqABC = rearrange(SeqABC, "t d jd -> 1 t d jd")
    SeqBCA = rearrange(SeqBCA, "t d jd -> 1 t d jd")
    SeqCAB = rearrange(SeqCAB, "t d jd -> 1 t d jd")

    B = np.concatenate([SeqABC, SeqBCA, SeqCAB], axis=0)

    return index_rearrange, B


def normalize_data(target_seq, left_seq, right_seq):
    """
    :param _seq: {n_frames x 57}
    """
    target_seq = target_seq.copy()
    left_seq = left_seq.copy()
    right_seq = right_seq.copy()
    if len(target_seq.shape) == 2:
        target_seq = rearrange(target_seq, "t (j d) -> t j d", j=19)
    if len(left_seq.shape) == 2:
        left_seq = rearrange(left_seq, "t (j d) -> t j d", j=19)
    if len(right_seq.shape) == 2:
        right_seq = rearrange(right_seq, "t (j d) -> t j d", j=19)
    target_seq_norm = reduce(target_seq[:, [6, 12]], "t j d -> j d", "mean")
    hip_left = target_seq_norm[0]
    hip_right = target_seq_norm[1]
    mu, R = get_normalization(left3d=hip_left, right3d=hip_right)
    target_seq = apply_normalization_to_seq(seq=target_seq, mu=mu, R=R)
    left_seq = apply_normalization_to_seq(seq=left_seq, mu=mu, R=R)
    right_seq = apply_normalization_to_seq(seq=right_seq, mu=mu, R=R)
    target_seq = rearrange(target_seq, "t j d -> t 1 (j d)")
    left_seq = rearrange(left_seq, "t j d -> t 1 (j d)")
    right_seq = rearrange(right_seq, "t j d -> t 1 (j d)")
    return np.concatenate([target_seq, left_seq, right_seq], axis=1)


class Residual(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        if dim_in == dim_out:
            self.proj = nn.Identity()
        else:
            self.proj = nn.Conv1d(dim_in, dim_out, kernel_size=1)
        self.nn = nn.Sequential(
            nn.Conv1d(
                in_channels=dim_in,
                out_channels=dim_out,
                kernel_size=3,
                padding=1,  # noqa E501
            ),
            nn.SiLU(),
            nn.BatchNorm1d(dim_out),
            nn.Conv1d(
                in_channels=dim_out,
                out_channels=dim_out,
                kernel_size=3,
                padding=1,  # noqa E501
            ),
        )

    def forward(self, x):
        """
        :param x: {n_batch x 57 x n_frames}
        """
        return self.nn(x) + self.proj(x)


class Down(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()

        self.nn = nn.Sequential(
            Residual(dim_in, dim_out),
            nn.SiLU(),
            nn.BatchNorm1d(dim_out),
            nn.Conv1d(dim_out, dim_out, kernel_size=3, padding=1, stride=2),
        )

    def forward(self, x):
        return self.nn(x)


class PositionalEncoding(nn.Module):
    def __init__(
        self, d_model: int, max_len: int = 2000, dropout: float = 0.1
    ):  # noqa E501
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Tensor, shape [seq_len, batch_size, embedding_dim]
            output: [sequence length, batch size, embed dim]
        """
        x = rearrange(x, "b d t -> t b d")
        x = x + self.pe[: x.size(0)]
        return rearrange(self.dropout(x), "t b d -> b t d")


class Model(nn.Module):
    def __init__(self, pose_dim, hidden_dim):
        super().__init__()
        self.down1_target = Down(pose_dim, hidden_dim)
        self.down1_left = Down(pose_dim, hidden_dim)
        self.down1_right = Down(pose_dim, hidden_dim)
        self.down2 = Down(hidden_dim * 3, hidden_dim * 4)
        self.down3 = Down(hidden_dim * 4, hidden_dim * 4)
        self.decision = nn.Sequential(
            Down(hidden_dim * 4, hidden_dim * 4),
            Down(hidden_dim * 4, hidden_dim * 4),
            Down(hidden_dim * 4, hidden_dim * 8),
            Down(hidden_dim * 8, hidden_dim * 8),
            nn.Conv1d(
                in_channels=hidden_dim * 8,
                out_channels=1,
                kernel_size=4,
                padding=0,  # noqa E501
            ),
        )

    def forward(self, Seq):
        """
        :param Seq: {n_batch x n_frames x 3 x 57}
        """
        target_seq = rearrange(Seq[:, :, 0], "b t d -> b d t")
        left_seq = rearrange(Seq[:, :, 1], "b t d -> b d t")
        right_seq = rearrange(Seq[:, :, 2], "b t d -> b d t")

        down1_target = self.down1_target(target_seq)
        down1_left = self.down1_left(left_seq)
        down1_right = self.down1_right(right_seq)
        down1 = torch.cat([down1_target, down1_left, down1_right], dim=1)

        down2 = self.down2(down1)
        down3 = self.down3(down2)

        dd = self.decision(down3)[:, :, 0]

        return dd


class WhoIsTheBuyerModule(pl.LightningModule):
    def __init__(self, pose_dim=57, hidden_dim=64):
        super().__init__()
        self.model = Model(pose_dim, hidden_dim)
        self.logit_bce_loss = nn.BCEWithLogitsLoss()

    def step(self, Seq, is_buyer):
        pred_logits = self.model(Seq)
        loss = self.logit_bce_loss(pred_logits, is_buyer)
        return pred_logits, loss

    def training_step(self, batch, batch_idx):
        Seq, is_buyer = batch
        is_buyer = is_buyer.unsqueeze(1)
        is_buyer_f = is_buyer.float()
        pred_logits, loss = self.step(Seq, is_buyer_f)
        with torch.no_grad():
            accuracy = torch.mean(
                (
                    (((F.sigmoid(pred_logits) > 0.5) * 1) == is_buyer) * 1
                ).float()  # noqa E501
            )

        self.log("acc", accuracy, on_step=True, on_epoch=False, prog_bar=True)

        return {"loss": loss, "acc": accuracy}

    def validation_step(self, batch, batch_idx):
        Seq, is_buyer = batch
        is_buyer = is_buyer.unsqueeze(1)
        is_buyer_f = is_buyer.float()
        pred_logits, loss = self.step(Seq, is_buyer_f)

        with torch.no_grad():
            accuracy = torch.mean(
                (
                    (((F.sigmoid(pred_logits) > 0.5) * 1) == is_buyer) * 1
                ).float()  # noqa E501
            )

        self.log(
            "val_loss", loss.item(), on_step=True, on_epoch=True, prog_bar=True
        )  # noqa E501
        self.log(
            "val_acc", accuracy, on_step=True, on_epoch=True, prog_bar=True
        )  # noqa E501
        return {"val_loss": loss.item(), "val_acc": accuracy}
social_diffusion/evaluation/buyer_decider.py
jutanke/social_diffusion
import hashlib
import math as m

from os.path import isfile
from typing import List

import numba as nb
import numpy as np

import numpy.linalg as la

# import torch
from annoy import AnnoyIndex

from social_diffusion.transforms.transforms import normalize

from einops import rearrange, reduce
from tqdm import tqdm


def compress_seq(seq, relevant_jids):
    """
    :param seq: {n_frames x 56}
    """
    seq = rearrange(seq.copy(), "t (j d) -> t j d", j=19, d=3)
    seq = normalize(seq, frame=0, jid_left=6, jid_right=12, fix_shape=False)
    return rearrange(seq[:, relevant_jids], "t j d -> (t j d)")


class Database:
    def __init__(self, scenes: List, motion_word_size=10, ss=1):
        self.relevant_jids = [1, 3, 4, 5, 7, 8, 9, 10, 11, 13, 14]

        self.compress_fn = compress_seq

        self.motion_word_size = motion_word_size
        assert int(m.ceil(motion_word_size / ss)) == int(
            m.floor(motion_word_size / ss)
        ), f"{m.ceil(motion_word_size / ss)} vs {int(m.floor(motion_word_size / ss))}"  # noqa E501
        n_dim = (len(self.relevant_jids) * 3) * motion_word_size // ss
        self.ss = ss
        self.meta = []
        self.lookup = AnnoyIndex(n_dim, "euclidean")

        # check if we have the db already!
        unique_txt = f"rel_ids: {self.relevant_jids} motion_word_size:{motion_word_size}, ss:{ss} ~ ~"  # noqa E501
        scene_names = sorted([scene.scene_name for scene in scenes])
        unique_txt += "++".join(scene_names)
        hash_object = hashlib.sha256(unique_txt.encode("utf-8"))
        unique_fname = f"/tmp/db_{hash_object.hexdigest()}.ann"

        if isfile(unique_fname):
            self.lookup.load(unique_fname)
        else:
            for scene in tqdm(scenes):
                poses = scene.get_poses()
                n_frames = poses.shape[0]
                n_person = poses.shape[1]
                speech = scene.get_speech()

                for t in range(n_frames - motion_word_size):
                    for i in range(n_person):
                        seq = poses[t : t + motion_word_size, i]  # noqa E501
                        speech_seq = speech[t : t + motion_word_size, i]  # noqa E501
                        try:
                            word = self.compress_fn(
                                seq[::ss], self.relevant_jids
                            )  # noqa E501
                            self.lookup.add_item(len(self.meta), word)
                            self.meta.append(
                                {
                                    "pid": i,
                                    "t": t,
                                    "name": scene.scene_name,
                                    "speech": speech_seq,
                                }
                            )
                        except:  # noqa E722
                            seq = rearrange(seq, "t (j d) -> t j d", j=19, d=3)
                            print(
                                f"\nFAILED AT t:{t} for pid{i} @{scene.scene_name}"  # noqa E501
                            )  # noqa E501
                            print("\tz-left:", seq[:, 6, 2])
                            print("\tz-right:", seq[:, 12, 2])
                            raise
            self.lookup.build(-1)
            self.lookup.save(unique_fname)

    def query(self, subseq):
        """
        :param subseq: {n_dim}
        """
        i = self.lookup.get_nns_by_vector(
            vector=subseq, n=1, include_distances=False
        )  # noqa E501
        i = i[0]
        true_seq = self.lookup.get_item_vector(i)
        kernel_size = self.motion_word_size
        true_seq = np.reshape(true_seq, (kernel_size, -1))
        query_seq = np.reshape(subseq, (kernel_size, -1))
        dist = ndms(true_seq, query_seq, kernel_size=kernel_size)
        return dist, i

    def rolling_query(self, subseq):
        """
        :param subseq: {n_frames x dim}
        """
        subseq = subseq.astype("float64")
        kernel_size = self.motion_word_size
        assert len(subseq) > kernel_size, str(subseq.shape)
        distances = []
        identities = []
        for frame in range((len(subseq) - kernel_size) + 1):
            sub_seq = subseq[frame : frame + kernel_size].copy()  # noqa E203
            word = self.compress_fn(sub_seq[:: self.ss], self.relevant_jids)
            word = word.flatten()
            d, i = self.query(word)
            distances.append(d)
            identities.append(i)
        distances = np.array(distances)
        return distances, identities

    def speech_inference(self, seq):
        """
        :param seq: {n_frames x 54}
        """
        n_frames = len(seq)
        assert n_frames >= self.motion_word_size

        speech = []
        for t in range(n_frames - self.motion_word_size):
            subseq = seq[t : t + self.motion_word_size]  # noqa E501
            try:
                word = self.compress_fn(subseq[:: self.ss], self.relevant_jids)
            except:  # noqa E722
                word = compress_seq(subseq, self.relevant_jids)
            i = self.lookup.get_nns_by_vector(
                vector=word, n=1, include_distances=False
            )[0]
            speech.append(np.mean(self.meta[i]["speech"]))
        speech = np.array(speech)
        return speech


@nb.njit(nb.float64[:, :](nb.float64[:, :]), nogil=True)
def velocity(seq):
    """
    :param [n_frames x dim]
    """
    n_frames = len(seq)
    dim = seq.shape[1]
    V = np.empty((n_frames - 1, dim), dtype=np.float64)
    for t in range(n_frames - 1):
        a = seq[t]
        b = seq[t + 1]
        v = b - a
        V[t] = v
    return V


@nb.njit(nb.float64(nb.float64[:, :], nb.float64[:, :], nb.int64), nogil=True)
def ndms(true_seq, query_seq, kernel_size):
    """
    :param true_seq: [kernel_size x n_useful_joints*3]
    :param query_seq: [kernel_size x n_useful_joints*3]
    """
    eps = 0.0000001
    true_v_ = np.ascontiguousarray(velocity(true_seq))
    query_v_ = np.ascontiguousarray(velocity(query_seq))
    dim = true_v_.shape[1]
    n_features = dim // 3
    true_v = true_v_.reshape((kernel_size - 1, n_features, 3))
    query_v = query_v_.reshape((kernel_size - 1, n_features, 3))
    total_score = 0.0
    for t in range(kernel_size - 1):
        for jid in range(n_features):
            a = np.expand_dims(true_v[t, jid], axis=0)
            b_T = np.expand_dims(query_v[t, jid], axis=1)
            norm_a = max(la.norm(a), eps)
            norm_b = max(la.norm(b_T), eps)
            cos_sim = (a @ b_T) / (norm_a * norm_b)
            disp = min(norm_a, norm_b) / max(norm_a, norm_b)
            score = ((cos_sim + 1) * disp) / 2.0
            total_score += score[0, 0] / n_features
    return total_score / (kernel_size - 1)
social_diffusion/evaluation/db.py
jutanke/social_diffusion
import os
from os.path import join, isfile


import pytorch_lightning as pl
import torch
import torch.nn.functional as F

from social_diffusion.transforms.transforms import to_canconical_form

from torch import nn, optim
from einops import rearrange

# from unio import isfile as isfile_unio, open as open_unio, UnioConfig


class Residual(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        if dim_in == dim_out:
            self.proj = nn.Identity()
        else:
            self.proj = nn.Conv1d(dim_in, dim_out, kernel_size=1)
        self.nn = nn.Sequential(
            nn.Conv1d(
                in_channels=dim_in,
                out_channels=dim_out,
                kernel_size=3,
                padding=1,  # noqa E501
            ),
            nn.SiLU(),
            nn.BatchNorm1d(dim_out),
            nn.Conv1d(
                in_channels=dim_out,
                out_channels=dim_out,
                kernel_size=3,
                padding=1,  # noqa E501
            ),
        )

    def forward(self, x):
        """
        :param x: {n_batch x 57 x n_frames}
        """
        return self.nn(x) + self.proj(x)


class Model(nn.Module):
    def __init__(self, pose_dim, hidden_dim, dropout=0.1, num_layers=3):
        """ """
        super().__init__()
        self.pose_encoder = nn.Linear(pose_dim, hidden_dim)
        self.gru_single = nn.GRU(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            dropout=dropout,
            batch_first=True,
            bidirectional=True,
            num_layers=num_layers,
        )
        self.merge = nn.Linear(2 * hidden_dim, hidden_dim)
        self.merge2 = nn.Linear(4 * hidden_dim, hidden_dim)

        self.gru_with_hist = nn.GRU(
            input_size=hidden_dim * 3,
            hidden_size=hidden_dim * 2,
            dropout=dropout,
            batch_first=True,
            bidirectional=True,
            num_layers=num_layers,
        )

        self.gru_with_hist2 = nn.GRU(
            input_size=hidden_dim * 3,
            hidden_size=hidden_dim * 2,
            dropout=dropout,
            batch_first=True,
            bidirectional=True,
            num_layers=num_layers,
        )

        self.intermediate = nn.Linear(2 * hidden_dim, 1)
        self.intermediate2 = nn.Linear(4 * hidden_dim, 1)
        self.final = nn.Linear(4 * hidden_dim, 1)

    def forward(self, Seq):
        """
        :param Seq: {n_batch x n_frames x 3 x 57}
        """
        emb = self.pose_encoder(Seq)

        seq0 = emb[:, :, 0]
        seq1 = emb[:, :, 1]
        seq2 = emb[:, :, 2]

        o0, _ = self.gru_single(seq0)
        o1, _ = self.gru_single(seq1)
        o2, _ = self.gru_single(seq2)

        out_interm0 = self.intermediate(o0)[:, :, 0]
        out_interm1 = self.intermediate(o1)[:, :, 0]
        out_interm2 = self.intermediate(o2)[:, :, 0]

        temp0 = self.merge(o0)
        temp1 = self.merge(o1)
        temp2 = self.merge(o2)

        combined_history = (
            temp0 + temp1 + temp2
        )  # order-invariant!  # TODO: RELU?  # noqa E501

        # TODO replace seq0 with temp0
        o0, _ = self.gru_with_hist(
            torch.cat([seq0, temp0, combined_history], dim=2)
        )  # noqa E501
        o1, _ = self.gru_with_hist(
            torch.cat([seq1, temp1, combined_history], dim=2)
        )  # noqa E501
        o2, _ = self.gru_with_hist(
            torch.cat([seq2, temp2, combined_history], dim=2)
        )  # noqa E501

        out_interm2_0 = self.intermediate2(o0)[:, :, 0]
        out_interm2_1 = self.intermediate2(o1)[:, :, 0]
        out_interm2_2 = self.intermediate2(o2)[:, :, 0]

        temp2_0 = self.merge2(o0)
        temp2_1 = self.merge2(o1)
        temp2_2 = self.merge2(o2)
        combined_history2 = temp2_0 + temp2_1 + temp2_2

        o0, _ = self.gru_with_hist2(
            torch.cat([seq0, temp2_0, combined_history2], dim=2)
        )
        o1, _ = self.gru_with_hist2(
            torch.cat([seq1, temp2_1, combined_history2], dim=2)
        )
        o2, _ = self.gru_with_hist2(
            torch.cat([seq2, temp2_2, combined_history2], dim=2)
        )

        out_final0 = self.final(o0)[:, :, 0]
        out_final1 = self.final(o1)[:, :, 0]
        out_final2 = self.final(o2)[:, :, 0]

        return (
            out_final0,
            out_final1,
            out_final2,
            out_interm0,
            out_interm1,
            out_interm2,
            out_interm2_0,
            out_interm2_1,
            out_interm2_2,
        )


class WhoIsTheBuyerModule(pl.LightningModule):
    def __init__(self, pose_dim=57, hidden_dim=128):
        super().__init__()
        self.model = Model(pose_dim, hidden_dim)
        self.logit_bce_loss = nn.BCEWithLogitsLoss()

    def step(self, Seq, speech):
        """
        :param Seq: {n_batch x n_frames x 3 x 57}
        :param speech: {n_batch x n_frames x 3}
        """
        speech0 = speech[:, :, 0]
        speech1 = speech[:, :, 1]
        speech2 = speech[:, :, 2]

        speech0_f = speech0.float()
        speech1_f = speech1.float()
        speech2_f = speech2.float()

        (
            out_final0,
            out_final1,
            out_final2,
            out_interm0,
            out_interm1,
            out_interm2,
            out_interm2_0,
            out_interm2_1,
            out_interm2_2,
        ) = self.model(Seq)

        loss0 = (
            self.logit_bce_loss(out_final0, speech0_f)
            + self.logit_bce_loss(out_interm0, speech0_f)
            + +self.logit_bce_loss(out_interm2_0, speech0_f)
        )
        loss1 = (
            self.logit_bce_loss(out_final1, speech1_f)
            + self.logit_bce_loss(out_interm1, speech1_f)
            + self.logit_bce_loss(out_interm2_1, speech1_f)
        )
        loss2 = (
            self.logit_bce_loss(out_final2, speech2_f)
            + self.logit_bce_loss(out_interm2, speech2_f)
            + self.logit_bce_loss(out_interm2_2, speech2_f)
        )

        with torch.no_grad():
            acc_interm0 = torch.mean(
                ((((F.sigmoid(out_interm0) > 0.5) * 1) == speech0) * 1).float()
            )
            acc_final0 = torch.mean(
                ((((F.sigmoid(out_final0) > 0.5) * 1) == speech0) * 1).float()
            )

            acc_interm1 = torch.mean(
                ((((F.sigmoid(out_interm1) > 0.5) * 1) == speech1) * 1).float()
            )
            acc_final1 = torch.mean(
                ((((F.sigmoid(out_final1) > 0.5) * 1) == speech1) * 1).float()
            )

            acc_interm2 = torch.mean(
                ((((F.sigmoid(out_interm2) > 0.5) * 1) == speech2) * 1).float()
            )
            acc_final2 = torch.mean(
                ((((F.sigmoid(out_final2) > 0.5) * 1) == speech2) * 1).float()
            )

            acc_final = (acc_final0 + acc_final1 + acc_final2) / 3.0
            acc_interm = (acc_interm0 + acc_interm1 + acc_interm2) / 3.0

        loss = loss0 + loss1 + loss2

        return loss, acc_final, acc_interm

    def training_step(self, batch, batch_idx):
        Seq, speech = batch
        loss, acc_final, acc_interm = self.step(Seq, speech)
        self.log(
            "acc_final", acc_final, on_step=True, on_epoch=True, prog_bar=True
        )  # noqa E501
        self.log(
            "acc_interm",
            acc_interm,
            on_step=True,
            on_epoch=True,
            prog_bar=True,  # noqa E501
        )  # noqa E501
        return {"loss": loss, "acc_final": acc_final, "acc_interm": acc_interm}

    def validation_step(self, batch, batch_idx):
        Seq, speech = batch
        loss, acc_final, acc_interm = self.step(Seq, speech)
        self.log(
            "val_acc_final",
            acc_final,
            on_step=False,
            on_epoch=True,
            prog_bar=True,  # noqa E501
        )
        self.log(
            "val_acc_interm",
            acc_interm,
            on_step=False,
            on_epoch=True,
            prog_bar=True,  # noqa E501
        )
        return {
            "val_loss": loss,
            "val_acc_final": acc_final,
            "val_acc_interm": acc_interm,
        }

    def configure_optimizers(self):
        optimizer = optim.Adam(self.model.parameters(), lr=0.0001)
        return optimizer

    def eval_chunk(self, P):
        """
        :param P: {n_frames x 3 x 57}
        """

        # P_canon = np.concatenate([seq0, seq1, seq2], axis=1)
        # P_canon = torch.from_numpy(P_canon).unsqueeze(0).to(self.device)
        with torch.no_grad():
            P = torch.from_numpy(P).to(self.device)

            seq0 = rearrange(
                to_canconical_form(P[:, 0])[:, :57], "t jd -> t 1 jd"
            )  # noqa E501
            seq1 = rearrange(
                to_canconical_form(P[:, 1])[:, :57], "t jd -> t 1 jd"
            )  # noqa E501
            seq2 = rearrange(
                to_canconical_form(P[:, 2])[:, :57], "t jd -> t 1 jd"
            )  # noqa E501

            P_canon = torch.cat([seq0, seq1, seq2], dim=1).unsqueeze(0)

            self.model.eval()
            (
                out_final0,
                out_final1,
                out_final2,
                _,
                _,
                _,
                _,
                _,
                _,
            ) = self.model(P_canon)

            out_final0 = (F.sigmoid(out_final0) > 0.5) * 1
            out_final1 = (F.sigmoid(out_final1) > 0.5) * 1
            out_final2 = (F.sigmoid(out_final2) > 0.5) * 1

        return rearrange(
            torch.cat([out_final0, out_final1, out_final2], dim=0)
            .cpu()
            .numpy(),  # noqa E501
            "p t -> t p",
        )


class WhoSpeaksDecider:
    def __init__(self, device):
        fname_local = join(os.getcwd(), "models/whotalks.ckpt")
        assert isfile(fname_local)

        self.model = WhoIsTheBuyerModule.load_from_checkpoint(fname_local)
        self.model.eval().to(device)
        self.device = device

    def eval(self, P):
        return self.model.eval_chunk(P)
social_diffusion/evaluation/decider.py
jutanke/social_diffusion
import math as m
import numpy as np
import numpy.linalg as la
import torch
from einops import rearrange, reduce
from scipy.ndimage import gaussian_filter1d
from social_diffusion.transforms.transforms import (
    get_normalization,
    apply_normalization_to_seq,
)


def get_hm(words, size=16):
    HM = np.zeros((size, size), dtype=np.float32)
    for t in range(len(words) - 1):
        a = words[t]
        b = words[t + 1]
        HM[a, b] += 1
    return HM


def to_prob(hm):
    hm = rearrange(hm, "w h -> (w h)")
    return hm / np.sum(hm)


def rot2d(theta):
    return np.array(
        [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]],
        dtype=np.float32,
    )


def get_normal(poses, left_jid: int, right_jid: int, smooth=True):
    """
    :param poses: {n_frames x 57}
    """
    is_torch = True
    if isinstance(poses, np.ndarray):
        is_torch = False
        poses = torch.from_numpy(poses)
    poses = rearrange(poses, "t (j d) -> t j d", j=19, d=3)
    left2d = poses[:, left_jid, :2]
    right2d = poses[:, right_jid, :2]
    y = right2d - left2d  # n_frames x 2
    y = y / (torch.norm(y, dim=1, p=2).unsqueeze(1) + 0.00000001)
    if not is_torch:
        y = y.numpy()
    # y[:, [0, 1]] = y[:, [1, 0]]

    R = rot2d(-m.pi / 2)

    y = y @ R

    if smooth:
        sigma = 5.0
        y[:, 0] = gaussian_filter1d(y[:, 0], sigma=sigma)
        y[:, 1] = gaussian_filter1d(y[:, 1], sigma=sigma)

    y = y / la.norm(y, axis=1, keepdims=True)

    return y


def calculate_head_orientations(poses):
    """
    :param poses: {n_frames x 3 x 57}
    :return
            {n_frames x n_person x 2 x 2}  # (pos / norm)
    """
    n_person = 3
    n_frames = poses.shape[0]
    head_orientations = np.empty((n_frames, n_person, 2, 2), dtype=np.float32)
    for i in range(n_person):
        seq = poses[:, i]
        head_orientations[:, i, 0] = rearrange(
            seq, "t (j d) -> t j d", j=19, d=3
        )[  # noqa E501
            :, 1, :2
        ]
        head_orientations[:, i, 1] = get_normal(seq, left_jid=15, right_jid=17)
    return head_orientations


def calculate_attn(P, buyer_id):
    """
    :param P: {n_frames x 3 x 57}
    """
    n_frames = P.shape[0]
    HO = calculate_head_orientations(P)
    seller_ids = [pid for pid in range(3) if pid != buyer_id]
    seller1 = HO[:, seller_ids[0]]
    seller2 = HO[:, seller_ids[1]]
    buyer = HO[:, buyer_id]

    seller1_loc = seller1[:, 0]
    seller2_loc = seller2[:, 0]
    buyer_loc = buyer[:, 0]
    buyer_lookdir = buyer[:, 1]

    def get_dir_vec(a, b):
        ab = b - a
        return ab / la.norm(ab, axis=1, keepdims=True)

    b_to_s1 = get_dir_vec(buyer_loc, seller1_loc)
    b_to_s2 = get_dir_vec(buyer_loc, seller2_loc)

    att1 = np.einsum("ij,ij->i", b_to_s1, buyer_lookdir)
    att2 = np.einsum("ij,ij->i", b_to_s2, buyer_lookdir)

    attn = np.zeros((n_frames, 2), dtype=np.int64)
    for t in range(n_frames):
        if att1[t] > att2[t]:
            attn[t, 0] = 1
        else:
            attn[t, 1] = 1
    return attn, seller_ids


def get_left_right_seller_ids(P, buyer_id):
    """
    :param P: {n_frames x 3 x 57}
    :return left id, right id
    """
    seller_ids = [pid for pid in range(3) if pid != buyer_id]
    buyer = P[:, buyer_id]
    seller1 = rearrange(P[:, seller_ids[0]], "t (j d) -> t j d", j=19, d=3)
    seller2 = rearrange(P[:, seller_ids[1]], "t (j d) -> t j d", j=19, d=3)

    buyer_mean_hip = reduce(
        buyer, "t (j d) -> j d", j=19, d=3, reduction="mean"
    )  # noqa E501

    mu, R = get_normalization(buyer_mean_hip[6], buyer_mean_hip[12])

    seller1_trans = apply_normalization_to_seq(seller1, mu=mu, R=R)
    seller2_trans = apply_normalization_to_seq(seller2, mu=mu, R=R)

    mean_seller1 = reduce(
        seller1_trans, "t j d -> d", j=19, d=3, reduction="mean"
    )  # noqa E501
    mean_seller2 = reduce(
        seller2_trans, "t j d -> d", j=19, d=3, reduction="mean"
    )  # noqa E501

    if mean_seller1[0] > mean_seller2[0]:
        return seller_ids[1], seller_ids[0]
    else:
        return seller_ids[0], seller_ids[1]


def attn_to_word5(attn, speech, buyer_speech):
    """
    :param attn: {n_frames x 2}
    :param speech: {n_frames x 2}
    :param buyer_speech: {n_frames}
    :param left_pid/right_pid  --> scaled to fit attn!!!
    """
    # 0 --> [0, l]    8 --> [b, l]
    # 1 --> [0, r]    9 --> [b, r]
    # 2 --> [L, l]   10 --> [bL, l]
    # 3 --> [L, r]   11 --> [bL, r]
    # 4 --> [R, l]   12 --> [bR, l]
    # 5 --> [R, r]   13 --> [bR, r]
    # 6 --> [LR, l]  14 --> [bLR, l]
    # 7 --> [LR, r]  15 --> [bLR, r]
    word_lookup_list = [
        ",l",
        ",r",
        "L,l",
        "L,r",
        "R,l",
        "R,r",
        "LR,l",
        "LR,r",
        "b,l",
        "b,r",
        "Lb,l",
        "Lb,r",
        "Rb,l",
        "Rb,r",
        "LRb,l",
        "LRb,r",
    ]
    assert len(word_lookup_list) == 16
    word_lookup = {}
    for i, word in enumerate(word_lookup_list):
        assert word not in word_lookup
        word_lookup[word] = i

    words = []
    for t in range(len(attn)):
        word = ""
        # -- step 1 -- handle who speaks
        if speech[t, 0] > 0.5:
            word += "L"
        if speech[t, 1] > 0.5:
            word += "R"
        if buyer_speech[t] > 0.5:
            word += "b"
        word += ","
        # -- step 2 -- handle attn
        if attn[t, 0] > 0.5:
            word += "l"
        elif attn[t, 1] > 0.5:
            word += "r"
        else:
            raise ValueError(f"wtf happend to the attention? @frame:{t}")
        words.append(word_lookup[word])
    return words


def words2segments(words, scene_name="None"):
    segments = []
    current_label = words[0]
    count = 1
    for i in range(1, len(words)):
        if words[i] == current_label:
            count += 1
        else:
            segments.append(
                {
                    "label": current_label,
                    "count": count,
                    "scene_name": scene_name,
                    "t": i - count,
                }
            )
            current_label = words[i]
            count = 1
    segments.append(
        {
            "label": current_label,
            "count": count,
            "scene_name": scene_name,
            "t": i - count,
        }
    )  # add final segment
    return segments
social_diffusion/evaluation/jsd.py
jutanke/social_diffusion
import math
import torch
import torch.nn as nn
from social_diffusion.transforms.rotation import tn_rot3d_c
from social_diffusion.diffusion import DiffusionModel
from social_diffusion.skeleton import Skeleton

from einops import rearrange, repeat, reduce
from einops.layers.torch import Rearrange


class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb


class TemporalConv1d(nn.Module):
    def __init__(self, dim_in, dim_out, kernel_size, stride):
        super().__init__()
        self.conv = nn.Conv1d(
            dim_in,
            dim_out,
            kernel_size=kernel_size,
            padding=(kernel_size - 1),
            stride=stride,
        )
        self.stride = stride
        self.kernel_size = kernel_size

    def forward(self, x):
        """
        param x: {n_batch x dim x n_frames}
        """
        chomp = self.kernel_size - self.stride
        return self.conv(x)[:, :, :-chomp].contiguous()


class ResidualBlock(nn.Module):
    def __init__(self, dim_in: int, dim_out: int, time_dim: int = -1):
        super().__init__()
        self.time_dim = time_dim
        if dim_in == dim_out:
            self.proj = nn.Identity()
        else:
            self.proj = nn.Conv1d(dim_in, dim_out, kernel_size=1)

        if self.time_dim > 0:
            self.mlp = nn.Sequential(
                nn.SiLU(), nn.Linear(time_dim, dim_out * 2)
            )  # noqa E501

        self.net = nn.Sequential(
            TemporalConv1d(
                dim_in=dim_in, dim_out=dim_out, kernel_size=3, stride=1
            ),  # noqa E501
            nn.SiLU(),
            Rearrange("b d t -> b t d"),
            nn.LayerNorm(dim_out),
            Rearrange("b t d -> b d t"),
            TemporalConv1d(
                dim_in=dim_out, dim_out=dim_out, kernel_size=3, stride=1
            ),  # noqa E501
        )

    def forward(self, x, time_emb=None):
        """
        :param x: {n_batch x dim x n_frames}
        :param time_emb: {n_batch x time_dim}
        """
        out = self.net(x)

        if self.time_dim > 0:
            time_scale, time_shift = self.mlp(time_emb).chunk(2, dim=1)
            time_scale = rearrange(time_scale, "b d -> b d 1")
            time_shift = rearrange(time_shift, "b d -> b d 1")
            out = out * (time_scale + 1) + time_shift

        return out + self.proj(x)


class Down(nn.Module):
    def __init__(self, dim_in, dim_out, time_dim=-1):
        super().__init__()
        self.block1 = ResidualBlock(
            dim_in=dim_out, dim_out=dim_out, time_dim=time_dim
        )  # noqa E501
        self.down = nn.Sequential(
            TemporalConv1d(
                dim_in=dim_in, dim_out=dim_out, kernel_size=3, stride=2
            ),  # noqa E501
        )

    def forward(self, x, time_embed=None):
        x = self.down(x)
        return self.block1(x, time_embed)


class Up(nn.Module):
    def __init__(self, dim_in, dim_out, time_dim=-1):
        super().__init__()
        self.block1 = ResidualBlock(
            dim_in=dim_out, dim_out=dim_out, time_dim=time_dim
        )  # noqa E501
        self.up = nn.Sequential(
            TemporalConv1d(
                dim_in=dim_in, dim_out=dim_out, kernel_size=3, stride=1
            ),  # noqa E501
            nn.Upsample(scale_factor=2, mode="linear"),
        )

    def forward(self, x, time_embed=None):
        x = self.up(x)
        return self.block1(x, time_embed)


def apply_normalization_per_person(Seq, R, T):
    """
    :param Seq: {n_batch x n_frames x n_person x d}
    :param R: {n_batch x n_person x 3 x 3}
    :param T: {n_batch x n_person x 3}
    """
    n_person = Seq.size(2)
    Seq_for_pid = []  # n_batch x n_person x n_person x d
    for pid in range(n_person):
        Seq_other = rearrange(Seq.clone(), "b t p d -> (b p) t d")
        R_pid = repeat(R[:, pid], "b d1 d2 -> (b p) d1 d2", p=n_person)
        T_pid = repeat(T[:, pid], "b d -> (b p) 1 d", p=n_person)
        Seq_other = Seq_other - T_pid
        Seq_other = torch.bmm(Seq_other, R_pid)
        Seq_other = rearrange(Seq_other, "(b p) t d -> b t p d", p=n_person)
        Seq_for_pid.append(Seq_other)
    Seq_for_pid = rearrange(
        torch.stack(Seq_for_pid), "p2 b t p1 d -> b t p1 p2 d"
    )  # noqa E501
    return Seq_for_pid


def undo_normalization_per_person(Seq, R, T):
    """
    :param Seq: {n_batch x n_frames x n_person x d}
    :param R: {n_batch x n_person x 3 x 3}
    :param T: {n_batch x n_person x 3}
    """
    n_person = Seq.size(2)
    Seq = rearrange(Seq, "b t p d -> (b p) t d")
    R_T = rearrange(R, "b p d1 d2 -> (b p) d2 d1")
    T = rearrange(T, "b p d -> (b p) 1 d")
    return rearrange(
        torch.bmm(Seq, R_T) + T, "(b p) t d -> b t p d", p=n_person
    )  # noqa E501


class MotionModel(DiffusionModel):
    def __init__(
        self,
        n_in,
        n_frames,
        skel: Skeleton,
        fourier_dim=64,
        hidden_dim=256,
        dropout=0.1,
    ):
        super().__init__(skel=skel)
        self.n_in = n_in
        pose_dim = self.jd()

        hidden_others = hidden_dim // 4

        time_dim = fourier_dim * 4
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(fourier_dim),
            nn.Linear(fourier_dim, time_dim),
            nn.GELU(),
            nn.Linear(time_dim, time_dim),
        )

        self.proj = nn.Conv1d(pose_dim, hidden_dim, kernel_size=1, padding=0)

        self.down1_others = Down(
            dim_in=pose_dim, dim_out=hidden_others, time_dim=time_dim
        )
        self.down2_others = Down(
            dim_in=hidden_others, dim_out=hidden_others, time_dim=time_dim
        )
        self.down3_others = Down(
            dim_in=hidden_others, dim_out=hidden_others, time_dim=time_dim
        )

        self.down1 = Down(
            dim_in=hidden_dim, dim_out=hidden_dim * 2, time_dim=time_dim
        )  # noqa E501
        self.down2 = Down(
            dim_in=hidden_dim * 2, dim_out=hidden_dim * 2, time_dim=time_dim
        )
        self.down3 = Down(
            dim_in=hidden_dim * 2, dim_out=hidden_dim * 2, time_dim=time_dim
        )
        self.res1 = ResidualBlock(
            dim_in=hidden_dim * 2, dim_out=hidden_dim * 2, time_dim=time_dim
        )

        self.up3 = Up(
            dim_in=hidden_dim * 2 + hidden_dim * 2,
            dim_out=hidden_dim * 2,
            time_dim=time_dim,
        )
        self.up2 = Up(
            dim_in=hidden_dim * 2 + hidden_dim * 2,
            dim_out=hidden_dim * 2,
            time_dim=time_dim,
        )
        self.up1 = Up(
            dim_in=hidden_dim * 2 + hidden_dim * 2,
            dim_out=hidden_dim,
            time_dim=time_dim,
        )

        self.final_res = ResidualBlock(
            dim_in=hidden_dim + hidden_dim,
            dim_out=hidden_dim,
            time_dim=time_dim,  # noqa E501
        )

        self.to_pose = nn.Sequential(
            nn.Conv1d(
                in_channels=hidden_dim,
                out_channels=pose_dim,
                kernel_size=1,
                padding=0,  # noqa E501
            ),
        )

        self.merge_history = nn.Conv1d(
            hidden_dim * 2 + hidden_others,
            hidden_dim * 2,
            kernel_size=1,
            padding=0,  # noqa E501
        )

    def forward(self, X_0, X_t, time, batch):
        """
        :param X_0: {n_batch x n_times x n_person x jd} clean gt image
        :param X_t: {n_batch x n_times x n_person x jd} properly noised
        :param batch: {
            "X_0": ...
            ...
        }
        """
        n_batch = X_t.size(0)
        n_frames = X_t.size(1)
        n_person = X_t.size(2)

        time = repeat(time, "b -> (b p)", p=n_person)
        t = self.time_mlp(time)

        # normalization:

        X_0_rot = X_0[:, :, :, 3:6]
        X_0_loc = X_0[:, :, :, :3]

        T = X_0_loc[:, self.n_in - 1, :, :].clone()
        X_0_rot = X_0_rot[:, self.n_in - 1, :, :]  # b p 3
        # print("newT", T.shape)
        # print("X_0_rot", X_0_rot.shape)

        x = X_0_rot[:, :, 0]
        y = X_0_rot[:, :, 1]
        rads = torch.atan2(y, x)
        rads = rearrange(rads, "b p -> (b p) 1")
        R = tn_rot3d_c(rads)
        R = rearrange(R, "(b p) d1 d2 -> b p d1 d2", p=n_person)

        X_t_pos = X_t[:, :, :, :3]

        X_t_rest = repeat(
            X_t[:, :, :, 3:], "b t p1 d -> b t p1 p2 d", p2=n_person
        )  # noqa E501

        X_t_pos = apply_normalization_per_person(X_t_pos, R, T)  # b t p1 p2 jd

        # rotate rotation points:
        #     * X_t_rot -> b t p 3
        #     * R -> b p 3 3
        X_t_rot = X_t_rest[:, :, :, :, 3:6]
        X_t_rest = X_t_rest[:, :, :, :, 3:]

        R_rot = repeat(R, "b p d1 d2 -> (b p p2) d1 d2", p2=n_person)
        X_t_rot = rearrange(X_t_rot, "b t p1 p2 d -> (b p1 p2) t d")

        X_t_rot = torch.bmm(X_t_rot, R_rot)
        X_t_rot = rearrange(
            X_t_rot,
            "(b p1 p2) t d -> b t p1 p2 d",
            p1=n_person,
            p2=n_person,  # noqa E501
        )
        X_t_all = torch.cat([X_t_pos, X_t_rot, X_t_rest], dim=4)

        X_t = []
        Hist = []
        for pid in range(n_person):
            X_t.append(X_t_all[:, :, pid, pid])

            x = rearrange(X_t_all[:, :, pid], "b t p jd -> (b p) jd t")
            d1 = self.down1_others(x, t)  # n_frames // 2
            d2 = self.down2_others(d1, t)  # n_frames // 4
            d3 = self.down3_others(d2, t)  # n_frames // 8
            d3 = reduce(d3, "(b p) d t -> b d t", "mean", p=n_person)
            Hist.append(d3)

        # 2, 64, 128, 32
        Hist = torch.stack(Hist)  # p2 b d t
        Hist = rearrange(Hist, "p b d t -> (b p) d t")

        # Hist = rearrange(torch.stack(Hist), )
        X_t = rearrange(torch.stack(X_t), "p b t jd -> b t p jd")

        X_t = rearrange(X_t, "b t p jd -> (b p) jd t")

        x = self.proj(X_t)

        d1 = self.down1(x, t)  # n_frames // 2
        d2 = self.down2(d1, t)  # n_frames // 4
        d3 = self.down3(d2, t)  # n_frames // 8

        h1 = self.res1(d3, t)  # n_frames // 8

        h2 = self.merge_history(torch.cat([h1, Hist], dim=1))

        u3 = self.up3(torch.cat([h2, d3], dim=1), t)  # n_frames // 4
        u2 = self.up2(torch.cat([u3, d2], dim=1), t)  # n_frames // 2
        u1 = self.up1(torch.cat([u2, d1], dim=1), t)  # n_frames

        final_inp = torch.cat([u1, x], dim=1)

        pose_out = self.to_pose(self.final_res(final_inp, t))

        pose_out = rearrange(
            pose_out,
            "(b p) jd t -> b t p jd",
            jd=self.jd(),
            p=n_person,
            t=n_frames,
            b=n_batch,
        )

        pose_out_loc = pose_out[:, :, :, :3]
        pose_out_rot = pose_out[:, :, :, 3:6]
        pose_out_rest = pose_out[:, :, :, 6:]

        pose_out_loc = undo_normalization_per_person(pose_out_loc, R=R, T=T)

        R_T = rearrange(R, "b p d1 d2 -> (b p) d2 d1")
        pose_out_rot = rearrange(pose_out_rot, "b t p d -> (b p) t d")
        pose_out_rot = torch.bmm(pose_out_rot, R_T)
        pose_out_rot = rearrange(
            pose_out_rot, "(b p) t d -> b t p d", p=n_person
        )  # noqa E501

        pose_out = torch.cat(
            [pose_out_loc, pose_out_rot, pose_out_rest], dim=3
        )  # noqa E501

        return pose_out
social_diffusion/model.py
jutanke/social_diffusion
import hashlib
from multiprocessing.pool import ThreadPool
from einops import rearrange, reduce
from typing import List
import numpy as np
import numba as nb
from social_diffusion import WeirdShapeError
from social_diffusion.skeleton import Skeleton


class Sequence:
    def __init__(self, n_frames):
        self.n_frames = n_frames

    def get_uid(self) -> str:
        raise NotImplementedError()

    def __len__(self):
        return self.n_frames


def get_uid(seq):
    return seq.get_uid()


class MultiPersonSequence(Sequence):
    """
    Represents an entire continous sequence of potentially multiple persons
    A mask defines if a person is visible (1) or not (0)
    """

    @staticmethod
    def global_uid(sequences: List):
        with ThreadPool(processes=len(sequences)) as p:
            return (
                "Gmps_"
                + hashlib.sha1(  # noqa B324
                    "___".join(p.map(get_uid, sequences)).encode("utf-8")
                ).hexdigest()
            )

    def __init__(
        self, Seq: np.ndarray, skel: Skeleton, Mask: np.ndarray = None
    ):  # noqa E501
        """
        :param Seq: {n_frames x n_person x jd}
        :param Mask: {n_frames x n_person} in {0, 1}
        """
        self.__uid = None
        self.n_frames = Seq.shape[0]
        self.n_person = Seq.shape[1]
        if Mask is None:
            Mask = np.ones((self.n_frames, self.n_person), dtype=np.float32)
        else:
            if (
                len(Mask.shape) != 2
                or Mask.shape[0] != self.n_frames
                or Mask.shape[1] != self.n_person
            ):
                raise WeirdShapeError(f"weird mask shape: {Mask.shape}")

        if len(Seq.shape) == 4:
            if Mask.shape[3] != 3:
                raise WeirdShapeError(f"(1) weird sequence shape: {Seq.shape}")
            if Seq.shape[2] != skel.n_joints:
                raise WeirdShapeError(
                    f"(1) Seq shape ({Seq.shape}) and Skeleton (j={skel.n_joints}) do not fit"  # noqa E501
                )
            Seq = rearrange("t p j d -> t p (j d)")
        if len(Seq.shape) != 3:
            raise WeirdShapeError(f"(2) weird sequence shape: {Seq.shape}")
        if Seq.shape[2] != 3 * skel.n_joints:
            f"(2) Seq shape ({Seq.shape}) and Skeleton (j={skel.n_joints}) do not fit"  # noqa E501

        super().__init__(n_frames=self.n_frames)

        self.skel = skel

        if Seq.dtype != np.float32:
            Seq = Seq.astype(np.float32)
        if Mask.dtype != np.float32:
            Mask = Mask.astype(np.float32)
        self.Seq = Seq * rearrange(Mask, "t p -> t p 1")
        self.Mask = Mask

    def flip_lr(self):
        """
        flips left-right jids
        """
        Seq_new = self.Seq.copy()
        Mask_new = self.Mask.copy()
        skel = self.skel
        for pid in range(self.n_person):
            Seq_new[:, pid] = skel.flip_lr(Seq_new[:, pid])
        return MultiPersonSequence(Seq=Seq_new, Mask=Mask_new, skel=skel)

    def get_uid(self) -> str:
        """
        To uniquely identify this sequence we produce a unique id
        ...quite hacky but gets the job done...
        """
        if self.__uid is None:
            uid_txt = f"n_frames:{self.n_frames}, n_person:{self.n_person},"
            uid_txt += f" mask:{np.sum(self.Mask)}"
            Seq_flat, Frames, Pids = nb_get_unmasked_sequences_flat(
                self.Seq, self.Mask
            )  # noqa E501
            Frames = np.array2string(Frames, threshold=len(Frames) + 1)
            Pids = np.array2string(Pids, threshold=len(Frames) + 1)
            uid_txt += f"ff:{Frames}==="
            uid_txt += f"pp:{Pids}"
            uid_txt += sequence_stats_to_unique_str(Seq_flat)
            self.__uid = (
                "mps" + hashlib.sha1(uid_txt.encode("utf-8")).hexdigest()
            )  # noqa B324
        return self.__uid

    def is_masked(self):
        """
        :returns True iff at least one frame has a masked pose
        """
        return np.min(self.Mask) < 0.9

    def n_max_person(self):
        return self.n_person

    def drop_empty_rows(self):
        M = (reduce(self.Mask, "t p -> p", "sum") > 0.5).nonzero()[0]
        if len(M) > 0 and len(M) < self.n_max_person():
            Seq = np.ascontiguousarray(self.Seq[:, M])
            Mask = np.ascontiguousarray(self.Mask[:, M])
            return MultiPersonSequence(Seq=Seq, skel=self.skel, Mask=Mask)
        else:
            return self

    def split_into_sequences_with_equal_n_person(self):
        """
        slits the sequence into subsquences with equal number of
        person each
        """
        all_sequences = []
        last_start = 0
        for t in range(1, self.n_frames):
            if np.sum(np.abs(self.Mask[t] - self.Mask[t - 1])) > 0.01:
                cur_Seq = self.Seq[last_start:t].copy()
                cur_Mask = self.Mask[last_start:t].copy()
                all_sequences.append(
                    MultiPersonSequence(
                        Seq=cur_Seq, Mask=cur_Mask, skel=self.skel
                    ).drop_empty_rows()
                )
                last_start = t
        # handle the last entry
        cur_Seq = self.Seq[last_start : t + 1].copy()  # noqa E501
        cur_Mask = self.Mask[last_start : t + 1].copy()  # noqa E501
        all_sequences.append(
            MultiPersonSequence(
                Seq=cur_Seq, Mask=cur_Mask, skel=self.skel
            ).drop_empty_rows()
        )
        return all_sequences


@nb.njit(
    nb.types.Tuple((nb.float32[:, :], nb.int64[:], nb.int64[:]))(
        nb.float32[:, :, :], nb.float32[:, :]
    ),
    nogil=True,
)
def nb_get_unmasked_sequences_flat(Seq, Mask):
    """
    :param Seq: {n_frames x n_person x jd}
    :param Mask: {n_frames x n_person}
    """
    n_frames = Seq.shape[0]
    n_person = Seq.shape[1]
    jd = Seq.shape[2]
    flat_Seq = np.empty((n_frames * n_person, jd), dtype=np.float32)
    Frames = np.empty((n_frames * n_person), dtype=np.int64)
    Pids = np.empty((n_frames * n_person), dtype=np.int64)

    ptr = 0
    for t in range(n_frames):
        for pid in range(n_person):
            if Mask[t, pid] > 0.5:
                flat_Seq[ptr] = Seq[t, pid]
                Frames[ptr] = t
                Pids[ptr] = pid
                ptr += 1

    return (
        np.ascontiguousarray(flat_Seq[:ptr]),
        np.ascontiguousarray(Frames[:ptr]),
        np.ascontiguousarray(Pids[:ptr]),
    )


def sequence_stats_to_unique_str(seq):
    """
    :param seq: {n_frames x jd}
    """
    if len(seq.shape) != 2:
        raise WeirdShapeError(f"Weird shape: {seq.shape}")
    ma = np.array2string(
        np.round(np.max(seq, axis=0), decimals=2), suppress_small=True
    )  # noqa E501
    mi = np.array2string(
        np.round(np.min(seq, axis=0), decimals=2), suppress_small=True
    )  # noqa E501
    mu = np.array2string(
        np.round(np.mean(seq, axis=0), decimals=2), suppress_small=True
    )
    std = np.array2string(
        np.round(np.std(seq, axis=0), decimals=2), suppress_small=True
    )
    return f"max:{ma}_min:{mi}_mu:{mu}_std:{std}"
social_diffusion/sequence.py
jutanke/social_diffusion
import numpy as np
from typing import List, Tuple
from social_diffusion import WeirdShapeError
from einops import rearrange
import torch

from social_diffusion.transforms.transforms import (
    matrix_to_axis_angle,
    tn_batch_normalize,
    tn_batch_rvec_to_rotpoint,
    tn_batch_rotpoint_to_rvec,
    tn_framewise_undo_transform,
)
from social_diffusion.transforms.transforms import axis_angle_to_matrix


class AbstractSkeleton:
    def __init__(
        self,
        is_local: bool,
        loss_weight: np.ndarray = None,
        leg_joints: List[int] = None,
    ):
        self.is_local = is_local
        if loss_weight is None:
            self.loss_weight = np.ones((self.jd()), dtype=np.float32)
        else:
            if len(loss_weight.shape) != 1 and len(loss_weight) != self.jd():
                raise WeirdShapeError(f"Weird loss weights: {loss_weight.shape}")
            self.loss_weight = loss_weight

        if leg_joints is not None:
            leg_joints = np.array(leg_joints, dtype=np.int64)
            if len(leg_joints.shape) != 1:
                raise WeirdShapeError(
                    f"Weird leg-joints shape: {leg_joints.shape}, req. 1d"
                )
        self.leg_joints = leg_joints

    def __len__(self):
        raise NotImplementedError("Nope")

    def get_leg_joints(self):
        if self.leg_joints is None:
            raise NotImplementedError("Nope...")
        else:
            return self.leg_joints

    def get_loss_weight(self):
        return self.loss_weight

    def jd(self):
        raise NotImplementedError("Nope")

    def from_local(self, device=None, use_rotpoints=True):
        raise NotImplementedError("Nope")

    def to_local(self, seq, device=None, is_flat=None, use_rotpoints=True):
        raise NotImplementedError("Nope")

    def jd_without_normalization(self):
        return self.jd()

    def from_local_group(self, Seq, device=None, use_rotpoints=True):
        """
        :param Seq: {n_frames x n_person x 6+jd}
        """
        n_person = Seq.shape[1]
        Seq_rec = []
        for pid in range(n_person):
            Seq_rec.append(
                self.from_local(Seq[:, pid], device=device, use_rotpoints=use_rotpoints)
            )
        return rearrange(Seq_rec, "p t jd -> t p jd")

    def to_local_group(self, Seq, device=None, is_flat=None, use_rotpoints=True):
        """
        :param Seq: {n_frames x n_person x 6+jd}
        """
        n_person = Seq.shape[1]
        Seq_loc = []
        for pid in range(n_person):
            Seq_loc.append(
                self.to_local(
                    Seq[:, pid],
                    device=device,
                    use_rotpoints=use_rotpoints,
                    is_flat=is_flat,
                )
            )
        return rearrange(Seq_loc, "p t jd -> t p jd")

    def fix_shape(self, seq, unroll_jd=False):
        """
        :param: seq: {n_frames x n_joints * 3}
        """
        n_joints = len(self)
        if len(seq.shape) == 1:
            if seq.shape[0] != self.jd():
                raise WeirdShapeError(f"(a) Weird shape: {seq.shape}")
            seq = rearrange(seq, "jd -> 1 jd")
        elif len(seq.shape) == 2:
            if seq.shape[1] == 3:
                if seq.shape[0] != n_joints:
                    raise WeirdShapeError(f"(a/1) Weird shape: {seq.shape}")
                seq = rearrange(seq, "j d -> 1 (j d)")
            elif seq.shape[1] == self.jd():
                seq = seq
            else:
                raise WeirdShapeError(f"(b) Weird shape: {seq.shape}, jd:{self.jd()}")
        elif len(seq.shape) == 3:
            if seq.shape[2] == 3:
                if seq.shape[1] != n_joints:
                    raise WeirdShapeError(
                        f"seq:{seq.shape} does not fit n_joints:{n_joints}"
                    )
                seq = rearrange(seq, "t j d -> t (j d)")
            else:
                raise WeirdShapeError(f"(c) Weird shape: {seq.shape}")
        if unroll_jd:
            return rearrange(seq, "t (j d) -> t j d", j=n_joints)
        else:
            return seq

    def split_local(self, seq):
        raise NotImplementedError("Nope")


class Skeleton(AbstractSkeleton):
    def __init__(
        self,
        normalizing_left_jid: int,
        normalizing_right_jid: int,
        n_joints: int,
        skeleton: List[Tuple[int, int]],
        left_jids: List[int],
        right_jids: List[int],
        *,
        standardization_mask: np.ndarray = None,
        loss_weight: np.ndarray = None,
        leg_joints: List[int] = None,
    ):
        self.normalizing_left_jid = normalizing_left_jid
        self.normalizing_right_jid = normalizing_right_jid
        self.n_joints = n_joints
        self.skeleton = skeleton
        self.left_jids = left_jids
        self.right_jids = right_jids
        self.standardization_mask = (
            standardization_mask
            if standardization_mask is not None
            else np.ones((1, self.jd()), dtype=np.float32)
        )
        super().__init__(is_local=False, loss_weight=loss_weight, leg_joints=leg_joints)

    def __len__(self):
        return self.n_joints

    def jd(self):
        return self.n_joints * 3

    def flip_lr(self, seq):
        """
        :param: seq: {n_frames x n_joints * 3}
        """
        seq = self.fix_shape(seq, unroll_jd=True)
        LR = np.array(list(self.left_jids) + list(self.right_jids), dtype=np.int64)
        RL = np.array(list(self.right_jids) + list(self.left_jids), dtype=np.int64)
        if len(LR) != 2 * len(self.left_jids):
            raise WeirdShapeError(
                f"Weird shape: {len(LR)} vs {len(self.left_jids)} (left)"
            )
        if len(RL) != 2 * len(self.right_jids):
            raise WeirdShapeError(
                f"Weird shape: {len(LR)} vs {len(self.right_jids)} (right)"
            )
        seq[:, LR] = seq[:, RL]
        I = np.diag([1.0, -1.0, 1.0])  # noqa: E741
        seq_flat = rearrange(seq, "t j d -> (t j) d")
        seq_flat = seq_flat @ I
        seq = rearrange(seq_flat, "(t j) d -> t j d", j=self.n_joints)
        return np.ascontiguousarray(rearrange(seq, "t j d -> t (j d)"))

    def to_local(self, seq, device=None, is_flat=None, use_rotpoints=True):
        """
        Poses are centered at the hip
        :param seq: {n_frames x jd}
        :param use_rotpoints: if True we use rotation points instead of rvec
            rotation points do not suffer from discontinuities
        :returns:
            {n_frames x 3 + 3 + jd}
            Mu + Rvec + jd
        """
        if is_flat is None:
            is_flat = len(seq.shape) == 2
        is_numpy = isinstance(seq, np.ndarray)
        seq = self.fix_shape(seq, unroll_jd=True)
        seq, R, Mu = tn_batch_normalize(
            Pose=seq, skel=self, return_transform=True, device=device
        )
        Rvec = matrix_to_axis_angle(R)
        if use_rotpoints:
            Rvec = tn_batch_rvec_to_rotpoint(Rvec, device=device)
        Rvec = rearrange(Rvec, "t d -> t 1 d")
        Mu = rearrange(Mu, "t d -> t 1 d")
        seq = torch.cat([Mu, Rvec, seq], dim=1)
        if is_numpy:
            seq = seq.cpu().numpy()
        if is_flat:
            seq = rearrange(seq, "t j d -> t (j d)")
        return seq

    def from_local(self, seq, device=None, use_rotpoints=True):
        """
        Poses are centered at the hip
        :param seq: {n_frames x jd}
        :returns:
            {n_frames x 3 + 3 + jd}
            Mu + Rvec + jd
        """
        is_flat = len(seq.shape) == 2
        is_numpy = isinstance(seq, np.ndarray)
        if is_numpy:
            seq = torch.from_numpy(seq)
        if device is None:
            device = torch.device("cpu")
        seq = seq.to(device)
        if is_flat:
            seq = rearrange(seq, "t (j d) -> t j d", d=3)
        if (
            len(seq.shape) != 3
            or seq.shape[1] != self.n_joints + 2
            or seq.shape[2] != 3
        ):
            raise WeirdShapeError(f"Weird shape: {seq.shape}, note that j=j+2")
        seq, Mu, Rvec = self.split_local(seq)
        if use_rotpoints:
            Rvec = tn_batch_rotpoint_to_rvec(Rvec, device)
        R = axis_angle_to_matrix(Rvec)
        seq = tn_framewise_undo_transform(poses=seq, mu=Mu, R=R)
        if is_numpy:
            seq = seq.cpu().numpy()
        if is_flat:
            seq = rearrange(seq, "b j d -> b (j d)")
        return seq

    def split_local(self, seq):
        """
        :param seq: {n_frames x jd}
        :returns:
            seq, rvec, mu
        """
        is_flat = False
        if len(seq.shape) == 2:
            seq = rearrange(seq, "t (j d) -> t j d", d=3)
            is_flat = True
        if (
            len(seq.shape) != 3
            or seq.shape[1] != self.n_joints + 2
            or seq.shape[2] != 3
        ):
            raise WeirdShapeError(
                f"Weird shape for splitting: {seq.shape}, J={self.n_joints}"
            )
        mu = seq[:, 0]
        rvec = seq[:, 1]
        poses = seq[:, 2:]
        if is_flat:
            poses = rearrange(poses, "t j d -> t (j d)")
        return poses, mu, rvec


class LocalSkeleton(AbstractSkeleton):
    """
    LocalSkeleton has the pose in canonical space and
    mu and rvec as first and second entry!
    [mu, rvec, pose]
    """

    def __init__(self, parent_skel: Skeleton, loss_weight: np.ndarray = None):
        if loss_weight is None:
            jd = parent_skel.jd() + 6
            loss_weight = np.ones((jd), dtype=np.float32)
            # loss_weight[:6] = 0.0
            # loss_weight[:3] = 1.0
            # loss_weight[3:] = 0

        self.normalizing_left_jid = parent_skel.normalizing_left_jid + 2
        self.normalizing_right_jid = parent_skel.normalizing_right_jid + 2
        self.n_joints = parent_skel.n_joints + 2

        leg_joints = None
        if parent_skel.leg_joints is not None:
            leg_joints = parent_skel.get_leg_joints() + 2

        self.parent_skel = parent_skel
        self.left_jids = (np.array(parent_skel.left_jids, dtype=np.int64) + 2).tolist()
        self.right_jids = (
            np.array(parent_skel.right_jids, dtype=np.int64) + 2
        ).tolist()
        self.skeleton = (np.array(parent_skel.skeleton, dtype=np.int64) + 2).tolist()
        self.standardization_mask = np.ones((1, self.jd()), dtype=np.float32)
        self.standardization_mask[0, :3] = 0
        super().__init__(is_local=True, loss_weight=loss_weight, leg_joints=leg_joints)

    def __len__(self):
        return self.n_joints

    def fix_shape(self, seq, unroll_jd=False):
        try:
            return super().fix_shape(seq, unroll_jd=unroll_jd)
        except WeirdShapeError:
            return self.parent_skel.fix_shape(seq, unroll_jd=unroll_jd)

    def jd_without_normalization(self):
        return self.jd() - 6

    def jd(self):
        return self.n_joints * 3

    def split_local(self, seq):
        return self.parent_skel.split_local(seq)

    def to_local(self, seq, *, use_rotpoints=True, is_flat=None, device=None):
        return self.parent_skel.to_local(
            seq, use_rotpoints=use_rotpoints, is_flat=is_flat, device=device
        )

    def from_local(self, seq, device=None, use_rotpoints=True):
        return self.parent_skel.from_local(
            seq, device=device, use_rotpoints=use_rotpoints
        )
social_diffusion/skeleton.py
jutanke/social_diffusion
import torch
from torch.optim import Adam
from torch.utils.data import DataLoader
from tqdm import tqdm
from social_diffusion import (
    create_training_dir,
    get_device,
    get_number_of_processes,
)  # noqa E501
from social_diffusion.diffusion import GaussianDiffusion

from social_diffusion.model import MotionModel
from social_diffusion.ema import EMA
from os.path import join, isdir

# from os import listdir
import numpy as np


def load_socdiff(
    model_dir: str,
    n_in: int,
    n_frames: int,
    hidden_dim: int,
    epoch: int,
    skel,
    n_person: int,
    device,
    *,
    objective="pred_x0",
):
    assert isdir(model_dir)
    fname_model = join(model_dir, f"model_{epoch}.pt")
    fname_ema = join(model_dir, f"ema_{epoch}.pt")

    model = MotionModel(
        n_in=n_in, n_frames=n_frames, skel=skel, hidden_dim=hidden_dim
    )  # noqa E501
    diffusion = GaussianDiffusion(
        model,
        n_person=n_person,
        n_in=n_in,
        n_frames=n_frames,
        objective=objective,
    ).to(device)
    ema_update_every = 10
    ema_decay = 0.995
    ema_diffusion = EMA(
        diffusion, beta=ema_decay, update_every=ema_update_every
    )  # noqa E501

    model.load_state_dict(torch.load(fname_model))
    ema_diffusion.load_state_dict(torch.load(fname_ema))

    return model, ema_diffusion


def train_socdiff(
    model_dir: str,
    ds_train,
    n_in,
    n_frames,
    n_person,
    *,
    batch_size=64,
    epochs=50,
    hidden_dim=256,
    device=None,
    tqdm_over_epochs=False,
    save_optimizer=False,
):
    """
    Train SocialDiffusion
    :param transform_batch_fn: def transform_batch_fn(batch) -> batch
    :param extra_loss_fn: def extra_loss_fn(batch, model_out, ds_train) ->
        loss (scalar)
    """
    train_lr = 1e-4

    print(f"Train SocDiff: {model_dir}")
    create_training_dir(model_dir)

    print(f"#train: {len(ds_train):,}")

    if device is None:
        device = get_device()
    objective = "pred_x0"

    model = MotionModel(
        n_in=n_in, n_frames=n_frames, skel=ds_train.skel, hidden_dim=hidden_dim
    )
    num_params = sum(
        param.numel() for param in model.parameters() if param.requires_grad
    )
    # logger.info(f"#params: {num_params:,} (hidden: {hidden_dim})")
    print(f"#params: {num_params:,} (hidden: {hidden_dim})")

    dl = DataLoader(
        ds_train,
        batch_size=batch_size,
        num_workers=get_number_of_processes(),
        shuffle=True,
    )

    adam_betas = (0.9, 0.99)
    diffusion = GaussianDiffusion(
        model,
        n_person=n_person,
        n_in=n_in,
        n_frames=n_frames,
        objective=objective,
    ).to(device)

    opt = Adam(diffusion.parameters(), lr=train_lr, betas=adam_betas)

    ema_update_every = 10
    ema_decay = 0.995
    ema_diffusion = EMA(
        diffusion, beta=ema_decay, update_every=ema_update_every
    )  # noqa E501

    for epoch in (
        pbar_epoch := tqdm(
            range(0, epochs),
            position=0,
            leave=True,
            disable=not tqdm_over_epochs,  # noqa E501
        )
    ):
        total_losses = []
        for batch in (
            pbar := tqdm(dl, position=0, leave=True, disable=tqdm_over_epochs)
        ):
            for k in batch.keys():
                batch[k] = batch[k].to(device)

            loss, model_out = diffusion(batch)

            loss.backward()

            # stupid pytorch bug...
            # torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

            opt.step()
            opt.zero_grad()

            total_losses.append(loss.item())

            if tqdm_over_epochs:
                pbar_epoch.set_description(
                    "epoch %04d loss %0.5f" % (epoch, np.mean(total_losses))
                )
            else:
                pbar.set_description(
                    "epoch %04d loss %0.5f" % (epoch, np.mean(total_losses))
                )

            ema_diffusion.update()

        fname_model = join(model_dir, f"model_{epoch}.pt")
        fname_ema = join(model_dir, f"ema_{epoch}.pt")
        fname_opt = join(model_dir, f"opt_{epoch}.pt")
        torch.save(model.state_dict(), fname_model)
        torch.save(ema_diffusion.state_dict(), fname_ema)
        if save_optimizer:
            torch.save(opt.state_dict(), fname_opt)

    ema_diffusion.eval()
    diffusion.eval()
    return {
        "model": model,
        "diffusion": diffusion,
        "ema_diffusion": ema_diffusion,
        "ds_train": ds_train,
        "n_in": n_in,
        "n_frames": n_frames,
    }
social_diffusion/train.py
jutanke/social_diffusion
import numpy as np
import torch


def tn_rot3d_c(c):
    """
    Rotation only around the z axis!
    :param c: {n_batch x 1}
    """
    if len(c.shape) == 1:
        c = c.unsqueeze(1)

    z = torch.zeros_like(c)
    o = torch.ones_like(c)

    cos_c = torch.cos(c)  # n_batch x 1
    sin_c = torch.sin(c)

    Rz_l1 = torch.cat([cos_c, -sin_c, z], dim=1).unsqueeze(1)
    Rz_l2 = torch.cat([sin_c, cos_c, z], dim=1).unsqueeze(1)
    Rz_l3 = torch.cat([z, z, o], dim=1).unsqueeze(1)

    Rz = torch.cat([Rz_l1, Rz_l2, Rz_l3], dim=1)

    return Rz


def rot3d(a, b, c):
    """"""
    Rx = np.array(
        [[1.0, 0.0, 0.0], [0, np.cos(a), -np.sin(a)], [0, np.sin(a), np.cos(a)]],
        np.float32,
    )
    Ry = np.array(
        [[np.cos(b), 0, np.sin(b)], [0.0, 1.0, 0.0], [-np.sin(b), 0, np.cos(b)]],
        np.float32,
    )
    Rz = np.array(
        [[np.cos(c), -np.sin(c), 0], [np.sin(c), np.cos(c), 0], [0.0, 0.0, 1.0]],
        np.float32,
    )
    return np.ascontiguousarray(Rx @ Ry @ Rz)


def apply_rotation_to_seq(seq, R):
    """
    :param seq: {n_frames x 18 x 3}
    :param R: {3x3}
    """
    R = np.expand_dims(R, axis=0)
    return np.ascontiguousarray(seq @ R)
social_diffusion/transforms/rotation.py
jutanke/social_diffusion
import numpy as np
import torch.nn.functional as F
from einops import rearrange
import torch
import numpy.linalg as la
import social_diffusion.transforms.rotation as rot
from social_diffusion import WeirdShapeError


class NormalizeDoesNotAllowZeroZError(ValueError):
    def __init__(self, msg):
        super().__init__(msg)


def apply_normalization_to_seq(seq, mu, R):
    """
    :param seq: {n_frames x 18 x 3}
    :param mu: {3}
    :param R: {3x3}
    """
    is_flat = False
    if len(seq.shape) == 2:
        is_flat = True
        seq = rearrange(seq, "t (j d) -> t j d", d=3)

    if len(seq.shape) != 3 or seq.shape[2] != 3:
        raise WeirdShapeError(f"weird shape: {seq.shape}")

    mu = np.expand_dims(np.expand_dims(np.squeeze(mu), axis=0), axis=0)
    seq = seq - mu
    seq_rot = rot.apply_rotation_to_seq(seq, R)

    if is_flat:
        seq_rot = rearrange(seq, "t j d -> t (j d)")
    return seq_rot


def get_normalization(left3d, right3d):
    """
    Get rotation + translation to center and face along the x-axis
    """
    mu = (left3d + right3d) / 2
    mu[2] = 0
    left2d = left3d[:2]
    right2d = right3d[:2]
    y = right2d - left2d
    y = y / (la.norm(y) + 0.00000001)
    angle = np.arctan2(y[1], y[0])
    R = rot.rot3d(0, 0, angle)
    return mu, R


def undo_normalization_to_seq(seq, mu, R):
    """
    :param seq: {n_frames x 18 x 3}
    :param mu: {3}
    :param R: {3x3}
    """
    is_flat = False
    if len(seq.shape) == 2:
        is_flat = True
        seq = rearrange(seq, "t (j d) -> t j d", d=3)
    mu = np.expand_dims(np.expand_dims(np.squeeze(mu), axis=0), axis=0)
    R_T = np.transpose(R)
    seq = rot.apply_rotation_to_seq(seq, R_T)
    seq = seq + mu
    if is_flat:
        seq = rearrange(seq, "t j d -> t (j d)")
    return seq


def tn_get_normalization(left3d, right3d, allow_zero_z=False):
    """
    :param left3d: {n_batch x 3}
    :param right3d: {n_batch x 3}
    Get rotation + translation to center and face along the x-axis
    """
    if not allow_zero_z:
        eps = 0.00000001
        left_z = eps > torch.abs(left3d[:, 2])
        right_z = eps > torch.abs(right3d[:, 2])
        if torch.count_nonzero(left_z) > 0 or torch.count_nonzero(right_z) > 0:
            print(
                f"left-> nonzero:{torch.count_nonzero(left_z)}, z:{left_z.shape}"  # noqa E501
            )  # noqa E501
            if torch.count_nonzero(left_z) > 0:
                for idx in torch.nonzero(left_z):
                    print(f"\t{idx} => {torch.round(left3d[idx], decimals=2)}")
                    if idx - 1 >= 0:
                        print(
                            f"\t\t t-1 => {torch.round(left3d[idx-1], decimals=2)}"  # noqa E501
                        )  # noqa E501
                    if idx + 1 < len(left3d):
                        print(
                            f"\t\t t+1 => {torch.round(left3d[idx+1], decimals=2)}"  # noqa E501
                        )  # noqa E501

            print(
                f"right-> nonzero:{torch.count_nonzero(right_z)}, z:{right_z.shape}"  # noqa E501
            )  # noqa E501
            if torch.count_nonzero(right_z) > 0:
                for idx in torch.nonzero(right_z):
                    print(f"\t{idx} => {right3d[idx]}")
            raise NormalizeDoesNotAllowZeroZError("We do not allow z:0!")

    mu = (left3d + right3d) / 2
    mu[:, 2] = 0

    left2d = left3d[:, :2]
    right2d = right3d[:, :2]
    y = right2d - left2d
    y = y / (torch.linalg.norm(y, dim=1, keepdims=True) + 0.0000001)
    a = y[:, 0]
    b = y[:, 1]
    angle = torch.atan2(b, a)

    R = rot.tn_rot3d_c(angle)
    return R, mu


def tn_batch_rotpoint_to_rvec(rotpoint, device=None):
    """ """
    is_numpy = False
    if isinstance(rotpoint, np.ndarray):
        is_numpy = True
        rotpoint = torch.from_numpy(rotpoint)
    if device is None:
        device = torch.device("cpu")
    if len(rotpoint.shape) != 2 or rotpoint.shape[1] != 3:
        raise ValueError(f"Weird rvec shape: {rotpoint.shape}")
    rotpoint = rotpoint.to(device)
    rotpoint = rotpoint / torch.norm(rotpoint, dim=1, keepdim=True)
    a = rotpoint[:, 0]
    b = rotpoint[:, 1]
    angle = torch.atan2(b, a)
    R = rot.tn_rot3d_c(angle)
    rvec = matrix_to_axis_angle(R)
    if is_numpy:
        rvec = rvec.cpu().numpy()
    return rvec


def tn_batch_rvec_to_rotpoint(rvec, device=None):
    """
    converts the rotation vector into a 3d point (that lives on the xy plane)
    The 3d point does not exhibit discontinuities!
    :param rvec: {n_batch x 3}
    """
    is_numpy = False
    if isinstance(rvec, np.ndarray):
        is_numpy = True
        rvec = torch.from_numpy(rvec)
    if device is None:
        device = torch.device("cpu")

    if len(rvec.shape) != 2 or rvec.shape[1] != 3:
        raise ValueError(f"Weird rvec shape: {rvec.shape}")
    rvec = rvec.to(device)
    Rs = axis_angle_to_matrix(rvec)
    pos2d = torch.zeros_like(rvec)
    pos2d[:, 0] = 1
    rotpoint = torch.bmm(Rs, pos2d.unsqueeze(2)).squeeze(2)
    if is_numpy:
        rotpoint = rotpoint.cpu().numpy()
    return rotpoint


def tn_batch_normalize(
    Pose: np.ndarray,
    skel,
    *,
    return_transform=False,
    allow_zero_z=False,
    device=None,  # noqa E501
):
    """
    :param Pose: {n_batch x jd}
    """
    Pose = skel.fix_shape(Pose, unroll_jd=True)
    if device is None:
        device = torch.device("cpu")
    if isinstance(Pose, np.ndarray):
        Pose = torch.from_numpy(Pose)
    if not torch.is_tensor(Pose):
        raise ValueError("<Pose> must be a torch Tensor!")

    left_jid = skel.normalizing_left_jid
    right_jid = skel.normalizing_right_jid

    Pose = Pose.to(device)

    Left3d = Pose[:, left_jid]
    Right3d = Pose[:, right_jid]

    from einops import reduce

    sum_left = (
        reduce(np.abs(Left3d.cpu().numpy()), "t d -> t", "sum") < 0.0001
    ) * 1  # noqa E501
    sum_right = (
        reduce(np.abs(Right3d.cpu().numpy()), "t d -> t", "sum") < 0.0001
    ) * 1  # noqa E501
    if np.count_nonzero(sum_left) > 0:
        print("left!!", np.count_nonzero(sum_left), sum_left.shape)
        print("right!!", np.count_nonzero(sum_right))
        exit()
    if np.count_nonzero(sum_right) > 0:
        print("right!!", np.count_nonzero(sum_right))
        exit()

    Rs, Mus = tn_get_normalization(
        left3d=Left3d, right3d=Right3d, allow_zero_z=allow_zero_z
    )
    Pose = tn_framewise_transform(poses=Pose, mu=Mus, R=Rs)

    if return_transform:
        return Pose, Rs, Mus
    else:
        return Pose


def tn_framewise_transform(poses, mu, R):
    """
    :param poses: {n_batch x j x 3}
    :param mu: {n_batch x 3}
    :param R: {n_batch x 3x3}
    """
    n_batch = poses.size(0)
    is_flat = False
    if len(poses.shape) == 2:
        is_flat = True
        poses = poses.reshape(n_batch, -1, 3)
    mu = mu.unsqueeze(1)
    poses = poses - mu
    output_poses = poses @ R
    if is_flat:
        output_poses = output_poses.reshape((n_batch, -1))
    return output_poses


def tn_framewise_undo_transform(poses, mu, R):
    """
    :param poses: {n_batch x j x 3}
    :param mu: {n_batch x 3}
    :param R: {n_batch x 3x3}
    """
    n_batch = poses.size(0)
    is_flat = False
    if len(poses.shape) == 2:
        is_flat = True
        poses = poses.reshape(n_batch, -1, 3)
    R_T = torch.transpose(R, 1, 2)
    mu = mu.unsqueeze(1)
    poses = poses @ R_T
    output_poses = poses + mu
    if is_flat:
        output_poses = output_poses.reshape((n_batch, -1))
    return output_poses


# ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~
# --- we do not want to have dependency on pytorch3d! ---
# instead we just graft the functions from their repo!
# https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py
# ALL CODE BELOW IS FROM PYTORCH3D!
# ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~ ~~


def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as axis/angle to rotation matrices.
    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))


def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to rotation matrices.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    r, i, j, k = torch.unbind(quaternions, -1)
    # pyre-fixme[58]: `/` is not supported for operand types `float` and
    # `Tensor`.
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return o.reshape(quaternions.shape[:-1] + (3, 3))


def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as axis/angle to quaternions.
    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles],
        dim=-1,  # noqa E501
    )
    return quaternions


def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as rotation matrices to axis/angle.
    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))


def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as rotation matrices to quaternions.
    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")

    batch_dim = matrix.shape[:-2]
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
        matrix.reshape(batch_dim + (9,)), dim=-1
    )

    q_abs = _sqrt_positive_part(
        torch.stack(
            [
                1.0 + m00 + m11 + m22,
                1.0 + m00 - m11 - m22,
                1.0 - m00 + m11 - m22,
                1.0 - m00 - m11 + m22,
            ],
            dim=-1,
        )
    )

    # we produce the desired quaternion multiplied by each of r, i, j, k
    quat_by_rijk = torch.stack(
        [
            # pyre-fixme[58]: `**` is not supported for operand types `Tensor`
            # and `int`.
            torch.stack(
                [q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1
            ),  # noqa E501
            # pyre-fixme[58]: `**` is not supported for operand types `Tensor`
            # and `int`.
            torch.stack(
                [m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1
            ),  # noqa E501
            # pyre-fixme[58]: `**` is not supported for operand types `Tensor`
            # and `int`.
            torch.stack(
                [m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1
            ),  # noqa E501
            # pyre-fixme[58]: `**` is not supported for operand types `Tensor`
            # and `int`.
            torch.stack(
                [m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1
            ),  # noqa E501
        ],
        dim=-2,
    )

    # We floor here at 0.1 but the exact level is not important; if q_abs is
    # small, the candidate won't be picked.
    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))

    # if not for numerical problems, quat_candidates[i] should be same (up to
    # a sign), forall i; we pick the best-conditioned one (with the largest
    # denominator)
    return quat_candidates[
        F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :
    ].reshape(batch_dim + (4,))


def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
    """
    Returns torch.sqrt(torch.max(0, x))
    but with a zero subgradient where x is 0.
    """
    ret = torch.zeros_like(x)
    positive_mask = x > 0
    ret[positive_mask] = torch.sqrt(x[positive_mask])
    return ret


def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to axis/angle.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles


# # ======== OLD ========
def normalize(
    seq,
    frame: int,
    jid_left=6,
    jid_right=12,
    return_transform=False,
    allow_zero_z=False,
    fix_shape=False,
):
    """
    :param seq: {n_frames x J x 3}
    :param frame:
    """
    if fix_shape:
        # seq = unknown_pose_shape_to_known_shape(seq)
        raise NotImplementedError("Not yet..")
    assert len(seq) > frame, f"frame:{frame} but shape: {seq.shape}"
    left3d = seq[frame, jid_left]
    right3d = seq[frame, jid_right]

    if not allow_zero_z:
        if np.isclose(left3d[2], 0.0):
            raise ValueError(f"Left seems to be zero! {left3d}")
        if np.isclose(right3d[2], 0.0):
            raise ValueError(f"Right seems to be zero! {right3d}")

    mu, R = get_normalization(left3d, right3d)
    if return_transform:
        return apply_normalization_to_seq(seq, mu, R), (mu, R)
    else:
        return apply_normalization_to_seq(seq, mu, R)


# def apply_normalization_to_seq(seq, mu, R, fix_shape=False):
#     """
#     :param seq: {n_frames x 18 x 3}
#     :param mu: {3}
#     :param R: {3x3}
#     """
#     mu = np.expand_dims(np.expand_dims(np.squeeze(mu), axis=0), axis=0)
#     if fix_shape:
#         raise NotImplementedError("nope")
#         # seq = unknown_pose_shape_to_known_shape(seq)
#     seq = seq - mu
#     return rot.apply_rotation_to_seq(seq, R, fix_shape=fix_shape)


# def unknown_pose_shape_to_known_shape(seq):  # noqa: C901
#     """
#     Sometimes we don't know what the pose is shaped like:
#     We could get either a single pose: (18*3) OR (18x3) OR (17*3) OR (17x3)
#     OR it could be a sequence: (nx18*3) OR (nx18x3) OR (nx17*3) OR (nx17x3)
#     This function takes any of those pose(s) and transforms
#     into (nx18x3) OR (nx17x3)
#     """
#     if isinstance(seq, list):
#         seq = np.array(seq, dtype=np.float32)
#     if len(seq.shape) == 1:
#         # has to be dimension 18*3 or 17*3
#         if seq.shape[0] == 18 * 3:
#             seq = seq.reshape(1, 18, 3)
#         else:
#             seq = seq.reshape(1, 17, 3)
#     elif len(seq.shape) == 2:
#         if seq.shape[0] == 18:
#             if seq.shape[1] != 3:
#                 raise ValueError("(1) Incorrect shape:" + str(seq.shape))
#             seq = seq.reshape(1, 18, 3)
#         elif seq.shape[0] == 17:
#             if seq.shape[1] != 3:
#                 raise ValueError("(1) Incorrect shape:" + str(seq.shape))
#             seq = seq.reshape(1, 17, 3)
#         else:
#             if seq.shape[1] != 18 * 3 and seq.shape[1] != 17 * 3:
#                 raise ValueError("(2) Incorrect shape:" + str(seq.shape))
#             n = len(seq)

#             if seq.shape[1] == 18 * 3:
#                 seq = seq.reshape(n, 18, 3)
#             else:
#                 seq = seq.reshape(n, 17, 3)
#     else:
#         if len(seq.shape) != 3:
#             raise ValueError("(3) Incorrect shape:" + str(seq.shape))
#         if (seq.shape[1] != 18 and seq.shape[1] != 17) or seq.shape[2] != 3:
#             raise ValueError("(4) Incorrect shape:" + str(seq.shape))
#     return seq


def to_canconical_form(seq):
    """
    :param seq: {n_frames x 57}
    :return: {n_frames x (57 + 2*3)}
    """
    was_numpy = False
    if isinstance(seq, np.ndarray):
        seq = torch.from_numpy(seq)
        was_numpy = True
    seq = rearrange(seq, "t (j d) -> t j d", j=19, d=3)
    R, mu = tn_get_normalization_for_pose(seq, left_jid=6, right_jid=12)
    seq_canonical = tn_framewise_transform(seq, mu, R)
    hips_world = seq[:, [6, 12]]
    seq_canonical = torch.cat([seq_canonical, hips_world], dim=1)
    if was_numpy:
        seq_canonical = seq_canonical.numpy()
    seq_canonical = rearrange(seq_canonical, "t j d -> t (j d)")
    return seq_canonical


def tn_get_normalization_for_pose(Pose, left_jid=13, right_jid=14):
    """
    0 1 2 3 4
    5 6 7 8 9 10
    11 12 [13 14] 15 16
    :param Pose: {n_batch x 17 x 3}
    """
    n_batch = len(Pose)
    if len(Pose.shape) == 2:
        # J = Pose.size(1) // 3
        Pose = Pose.reshape(n_batch, 17, 3)
    assert len(Pose.shape) == 3, f"wrong shape: {Pose.shape}"
    left3d = Pose[:, left_jid]
    right3d = Pose[:, right_jid]
    return tn_get_normalization(left3d=left3d, right3d=right3d)
social_diffusion/transforms/transforms.py
jutanke/social_diffusion
import shutil
from os import listdir, makedirs, remove
from os.path import isdir as isdir_local, isfile as isfile_local, join
import cv2


def to_mp4(path, delete_pngs=True, fps=25.0, nofolder=True):
    """
    creates a video file from the given images at the path
    """
    if nofolder:
        if not delete_pngs:
            raise ValueError("the folder with all pngs will be deleted!")
    if not isdir_local(path):
        raise ValueError(f"Path not exist: {path}")
    # fourcc = cv2.VideoWriter_fourcc(*"H264")
    # fourcc = cv2.VideoWriter_fourcc(*"avc1")  # avc1 == h264
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = None
    for fname in sorted(
        [join(path, f) for f in listdir(path) if f.endswith(".png")]
    ):  # noqa E501
        im = cv2.imread(fname)
        h = im.shape[0]
        w = im.shape[1]
        if out is None:
            out = cv2.VideoWriter(join(path, "out.mp4"), fourcc, fps, (h, w))
        out.write(im)
        if delete_pngs:
            remove(fname)
    out.release()

    if nofolder:
        root_path = "/".join(path.split("/")[:-1])
        fname_target = join(root_path, path.split("/")[-1] + ".mp4")
        fname_source = join(path, "out.mp4")
        shutil.move(fname_source, fname_target)
        # from time import sleep

        # time.sleep(1)
        shutil.rmtree(path)


def create_vis_path_local(path: str):
    """
    Creates the visualization path. If the path exists
    it is being deleted gracefully and then re-created
    """
    if isdir_local(path):
        # double check that in here we ONLY have
        # visualized data and NOTHING else!
        allowed_extensions = [".png", ".mp4", ".gif"]
        for fname in listdir(path):
            fname = join(path, fname)
            if isfile_local(fname):
                fname = fname.lower()
                is_ok = False
                for ext in allowed_extensions:
                    if fname.endswith(ext):
                        is_ok = True
                        continue
                if not is_ok:
                    raise ValueError(
                        f"<save_create_vis_path> path {path} contains forbidden files: {fname}!"  # noqa E501
                    )
        # folder is deemed save to delete
        shutil.rmtree(path)
    makedirs(path)

social_diffusion/vis/__init__.py
jutanke/social_diffusion
from os.path import join

import matplotlib.pylab as plt
import numpy as np
from einops import rearrange
from mpl_toolkits import mplot3d  # noqa E501 ensure that we can plot 3d
from social_diffusion.skeleton import Skeleton
from social_diffusion.vis import create_vis_path_local, to_mp4
from tqdm import tqdm


def plot_seq(
    seq,
    path: str,
    skel: Skeleton,
    *,
    predict_from_frame: int = -1,
    to_video: bool = True,
    figsize: int = 10,
    noaxis: bool = True,
    is_group: bool = False,
    alpha: float = 1.0,
    linewidth: float = 3.0,
    keep_past_frames: bool = False,
    plot_jids: bool = False,
    mark_origin=False,
):
    """
    :param seq: {n_frames x jd}
    """
    create_vis_path_local(path)
    if is_group:
        if len(seq.shape) == 3 or len(seq.shape) == 4:
            Seq = []
            n_person = seq.shape[1]
            for pid in range(n_person):
                Seq.append(skel.fix_shape(seq[:, pid], unroll_jd=True))
            seq = rearrange(
                np.array(Seq, dtype=np.float32), "p t j d -> t p j d"
            )  # noqa E501
        else:
            raise ValueError(f"Weird shape for group: {seq.shape}")
        seq_flat = rearrange(seq, "p t j d -> (p t) j d")
    else:
        seq = skel.fix_shape(seq, unroll_jd=True)
        seq_flat = seq
    n_frames = len(seq)

    if predict_from_frame == -1:
        predict_from_frame = n_frames + 1

    fig = plt.figure(figsize=(figsize, figsize))
    ax = fig.add_subplot(111, projection="3d")

    if keep_past_frames:
        alphas = np.linspace(0.00001, alpha, num=n_frames)
    else:
        alphas = [alpha] * n_frames

    for t in tqdm(range(n_frames), total=n_frames, leave=True, position=0):
        if not keep_past_frames:
            ax.clear()
            ax.set_title(f"frame {t}")
        if noaxis:
            ax.axis("off")

        if mark_origin:
            ax.plot([-1, 1], [0, 0], [0, 0], color="black", alpha=0.5)
            ax.plot([0, 0], [-1, 1], [0, 0], color="black", alpha=0.5)
            ax.plot([0, 0], [0, 0], [-1, 1], color="black", alpha=0.5)

        set_axis(ax, seq_flat, skel)
        if t >= predict_from_frame:
            lcolor = "orange"
            rcolor = "green"
        else:
            lcolor = "cornflowerblue"
            rcolor = "salmon"
        plot_pose(
            ax,
            seq[t],
            skel,
            left_color=lcolor,
            right_color=rcolor,
            alpha=alphas[t],
            linewidth=linewidth,
            plot_jids=plot_jids,
        )
        fname = join(path, "frame%05d.png" % t)
        plt.tight_layout()
        plt.savefig(fname)

    if to_video:
        to_mp4(path)


def set_axis(ax, pose, skel: Skeleton):
    """
    :param ax:
    :param pose: {n_person x jd}
    :param skel: Skel
    """
    pose = skel.fix_shape(pose, unroll_jd=True)  # p x j x d
    X = pose[:, :, 0]
    Y = pose[:, :, 1]
    Z = pose[:, :, 2]

    Xmin = np.min(X)
    Ymin = np.min(Y)
    Zmin = np.min(Z)
    Xmax = np.max(X)
    Ymax = np.max(Y)
    Zmax = np.max(Z)

    max_size = max(Xmax - Xmin, max(Ymax - Ymin, Zmax - Zmin))
    ax.set_xlim(
        [(Xmax + Xmin) / 2 - max_size / 2, (Xmax + Xmin) / 2 + max_size / 2]
    )  # noqa E501
    ax.set_ylim(
        [(Ymax + Ymin) / 2 - max_size / 2, (Ymax + Ymin) / 2 + max_size / 2]
    )  # noqa E501
    ax.set_zlim(
        [(Zmax + Zmin) / 2 - max_size / 2, (Zmax + Zmin) / 2 + max_size / 2]
    )  # noqa E501


def plot_pose(
    ax,
    pose,
    skel: Skeleton,
    *,
    left_color="cornflowerblue",
    right_color="salmon",
    center_color="gray",
    plot_jids=False,
    alpha=1.0,
    linewidth=3,
):
    """
    :param ax:
    :param pose: {n_person x jd}
    :param skel: Skel
    """
    pose = skel.fix_shape(pose, unroll_jd=True)  # p x j x d
    Alpha = None
    if isinstance(alpha, list) or isinstance(alpha, np.ndarray):
        Alpha = alpha
        if len(Alpha) != len(pose):
            raise ValueError(
                f"Alpha ({len(Alpha)}) and pose ({len(pose)}) have to have the same length"  # noqa E501
            )

    n = len(pose)
    if Alpha is None:
        Alpha = [alpha for _ in range(n)]
    for a, b in skel.skeleton:
        if a in skel.left_jids and b in skel.left_jids:
            color = left_color
        elif a in skel.right_jids and b in skel.right_jids:
            color = right_color
        else:
            color = center_color

        for t in range(n):
            px = pose[t, a, 0]
            py = pose[t, a, 1]
            pz = pose[t, a, 2]

            qx = pose[t, b, 0]
            qy = pose[t, b, 1]
            qz = pose[t, b, 2]

            ax.plot(
                [px, qx],
                [py, qy],
                [pz, qz],
                color=color,
                linewidth=linewidth,
                alpha=Alpha[t],
            )

    if plot_jids:
        for t in range(n):
            for jid, loc in enumerate(pose[t]):
                ax.text(loc[0], loc[1], loc[2], str(jid))

social_diffusion/vis/plot.py
jutanke/social_diffusion
from social_diffusion import get_output_dir
from os.path import join, isdir
from os import makedirs
from social_diffusion.datasets.haggling import get_haggling_test_sequences

# from social_diffusion.vis.plot import plot_seq
from social_diffusion.datasets.haggling import get_haggling_local_dataset
from social_diffusion.train import train_socdiff
import numpy as np


def train():
    print("start training on Haggling")

    n_in = 128
    n_frames = 512
    n_person = 3
    hidden_dim = 128 + 64

    # test_seqs = get_haggling_test_sequences()

    model_dir = join(
        get_output_dir(),
        f"models/socdiff_haggling_nin{n_in}_nframes{n_frames}_hidden{hidden_dim}",  # noqa E501
    )

    ds_train, _, skel = get_haggling_local_dataset(
        n_frames=n_frames, n_in=n_in
    )  # noqa E501

    training_data = train_socdiff(
        model_dir=model_dir,
        ds_train=ds_train,
        n_in=n_in,
        n_frames=n_frames,
        n_person=n_person,
        hidden_dim=hidden_dim,
    )

    test_seqs = get_haggling_test_sequences()

    OUT = []
    for test_seq in test_seqs:
        Seq = test_seq.Seq
        P = np.zeros((2048, 3, 57), dtype=np.float32)
        P[:128] = Seq[:128]
        ema_diffusion = training_data["ema_diffusion"]
        Seq_hat = ema_diffusion.ema_model.predict(
            P, ds_train=ds_train, n_samples=8
        )  # noqa E501
        OUT.append(Seq_hat)

    outpath = join(get_output_dir(), "results/haggling")
    if not isdir(outpath):
        makedirs(outpath)
    fname = join(outpath, "pred.npy")
    OUT = np.array(OUT, dtype=np.float32)
    np.save(fname, OUT)


if __name__ == "__main__":
    train()

train_haggling.py
jutanke/social_diffusion
from social_diffusion import get_output_dir
from social_diffusion.vis.plot import plot_seq
from social_diffusion.train import load_socdiff
from social_diffusion.datasets.panoptic import PANOPTIC_SKELETON
from os.path import join
import numpy as np
from multiprocessing import Pool
from social_diffusion.datasets.haggling import (
    get_haggling_local_dataset,
    get_haggling_test_sequences,
)
import torch
from einops import repeat


def render_seq(entry):
    Seq = entry["Seq"]
    vis_path = entry["vis_path"]
    plot_seq(
        seq=Seq,
        path=vis_path,
        skel=PANOPTIC_SKELETON,
        is_group=True,
        predict_from_frame=178,
    )


def vis():
    n_in = 128
    n_frames = 512
    n_person = 3
    hidden_dim = 128 + 64

    model_dir = join(
        get_output_dir(),
        f"models/socdiff_haggling_nin{n_in}_nframes{n_frames}_hidden{hidden_dim}",  # noqa E501
    )
    device = torch.device("cuda")

    ds_train, _, skel = get_haggling_local_dataset(
        n_frames=n_frames, n_in=n_in
    )  # noqa E501

    _, ema_diffusion = load_socdiff(
        model_dir=model_dir,
        n_in=n_in,
        n_frames=n_frames,
        hidden_dim=hidden_dim,
        epoch=50,
        skel=skel,
        n_person=n_person,
        device=device,
    )

    test_seqs = get_haggling_test_sequences()

    seq_idx = 0
    n_samples = 8
    JUMPOFF_FRAME = 178

    test_seq = test_seqs[seq_idx]
    Seq = test_seq.Seq
    P = np.zeros((2048, 3, 57), dtype=np.float32)
    P[:128] = Seq[JUMPOFF_FRAME - 128 : JUMPOFF_FRAME]  # noqa E203
    Seq_hat = ema_diffusion.ema_model.predict(
        P, ds_train=ds_train, n_samples=n_samples
    )  # noqa E501

    Seq_hat[:, :128] = Seq[JUMPOFF_FRAME - 128 : JUMPOFF_FRAME]  # noqa E203

    Seq_hat_init = repeat(
        Seq[: JUMPOFF_FRAME - 128], "t p jd -> s t p jd", s=n_samples
    )  # noqa E501

    Seq_hat = np.concatenate([Seq_hat_init, Seq_hat], axis=1)

    Data = []
    for i in range(8):
        Seq = Seq_hat[i, :512]
        vis_path = join(get_output_dir(), f"haggling{seq_idx}/vis{i}")
        Data.append({"Seq": Seq, "vis_path": vis_path})
    with Pool(8) as p:
        p.map(render_seq, Data)


if __name__ == "__main__":
    vis()

vis_haggling.py
jutanke/social_diffusion