eclm.py

import argparse
import os
import cv2
import torch
import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt
from tqdm import tqdm
from copy import deepcopy
from ultralytics import YOLO
from model.hmr2 import hmr2
from data.preprocess import camerahmr
from data.preprocess.camerahmr.cam_model.fl_net import FLNet
from data.preprocess.camerahmr.datasets.utils import generate_image_patch_cv2
from configs import constants as _C
from torchvision.transforms import Normalize
from smplx import SMPL

NUM_LAYERS = 32
PERSON_ID = 0
NUM_VIDEOS = 17


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_name', type=str, required=True)
    parser.add_argument('--n_samples', type=int, default=1360)
    parser.add_argument('--example_type', type=str,
                        choices=['iid', 'continuous'], default='iid')
    parser.add_argument('--type', type=str,
                        choices=['analysis', 'eclm'], required=True)
    parser.add_argument("--threshold", type=float, default=0.1)
    args = parser.parse_args()
    return args


def init_cam_model(device):
    model = FLNet()
    checkpoint = torch.load(_C.PATHS.CAMERAMODEL)['state_dict']
    model.load_state_dict(checkpoint)
    model = model.to(device)
    model.eval()
    return model


def init_model(model_name, type, device):
    hidden_states_indices = list(
        range(0, 32)) if type == 'analysis' else 'last'
    if model_name == 'camerahmr':
        model = camerahmr.CameraHMR.load_from_checkpoint(
            _C.PATHS.CAMERAHMR, strict=False, hidden_states_indices=hidden_states_indices)
    elif model_name == 'hmr2':
        model = hmr2(_C.PATHS.HMR2, hidden_states_indices)
    else:
        raise ValueError(f"Model {model_name} not supported.")
    model = model.to(device)
    model.eval()
    return model


def resize_image(img, target_size):
    height, width = img.shape[:2]
    aspect_ratio = width / height

    # Calculate the new size while maintaining the aspect ratio
    if aspect_ratio > 1:
        new_width = target_size
        new_height = int(target_size / aspect_ratio)
    else:
        new_width = int(target_size * aspect_ratio)
        new_height = target_size

    # Resize the image using OpenCV
    resized_img = cv2.resize(img, (new_width, new_height),
                             interpolation=cv2.INTER_AREA)

    # Create a new blank image with the target size
    final_img = np.ones((target_size, target_size, 3), dtype=np.uint8) * 255

    # Paste the resized image onto the blank image, centering it
    start_x = (target_size - new_width) // 2
    start_y = (target_size - new_height) // 2
    final_img[start_y:start_y + new_height,
              start_x:start_x + new_width] = resized_img

    return aspect_ratio, final_img


def expand_to_aspect_ratio(input_shape, target_aspect_ratio=None):
    """Increase the size of the bounding box to match the target shape."""
    if target_aspect_ratio is None:
        return input_shape

    try:
        w, h = input_shape
    except (ValueError, TypeError):
        return input_shape

    w_t, h_t = target_aspect_ratio
    if h / w < h_t / w_t:
        h_new = max(w * h_t / w_t, h)
        w_new = w
    else:
        h_new = h
        w_new = max(h * w_t / h_t, w)
    if h_new < h or w_new < w:
        breakpoint()
    return np.array([w_new, h_new])


def convert_cvimg_to_tensor(cvimg: np.array):
    """
    Convert image from HWC to CHW format.
    Args:
        cvimg (np.array): Image of shape (H, W, 3) as loaded by OpenCV.
    Returns:
        np.array: Output image of shape (3, H, W).
    """
    # from h,w,c(OpenCV) to c,h,w
    img = cvimg.copy()
    img = np.transpose(img, (2, 0, 1))
    # from int to float
    img = img.astype(np.float32)
    return img


def trans_points2d_parallel(keypoints_2d: np.array, trans: np.array):

    ones = np.ones((keypoints_2d.shape[0], 1))
    keypoints_augmented = np.hstack([keypoints_2d, ones])

    # Perform the transformation: (3, 3) dot (N, 3)^T => (3, N)
    transformed_keypoints = np.dot(trans, keypoints_augmented.T).T

    # Extract the transformed 2D points by discarding the last row
    transformed_keypoints_2d = transformed_keypoints[:, :2]

    return transformed_keypoints_2d


def compute_cka(hidden1, hidden2):
    """
    hidden1 and hidden2: (batch_size, num_tokens, num_channels)
    """
    K = hidden1 @ hidden1.permute(0, 2, 1)
    L = hidden2 @ hidden2.permute(0, 2, 1)

    b, n = hidden1.shape[:2]
    H = (
        torch.eye(n) - ((1 / n) * torch.ones((n, n)))
    ).unsqueeze(0).expand(b, -1, -1).to(hidden1.device)

    K_bar = H @ K @ H
    L_bar = H @ L @ H

    cka = torch.cosine_similarity(K_bar.flatten(1), L_bar.flatten(1))
    return cka.mean()


def get_example(img_path, center_x, center_y,
                width, height,
                keypoints_2d,
                patch_width, patch_height,
                mean, std,
                is_bgr=True,
                border_mode=cv2.BORDER_CONSTANT,
                do_flip=False):

    if isinstance(img_path, str):
        # 1. load image
        cvimg = cv2.imread(img_path, cv2.IMREAD_COLOR |
                           cv2.IMREAD_IGNORE_ORIENTATION)
        if not isinstance(cvimg, np.ndarray):
            raise IOError("Fail to read %s" % img_path)
    elif isinstance(img_path, np.ndarray):
        cvimg = img_path
    else:
        raise TypeError('img_path must be either a string or a numpy array')

    img_height, img_width, img_channels = cvimg.shape
    img_size = np.array([img_height, img_width])

    scale, rot, color_scale, tx, ty = 1.0, 0, [1.0, 1.0, 1.0], 0., 0.

    assert width >= 1 and height >= 1, f"width: {width}, height: {height}"

    center_x += width * tx
    center_y += height * ty

    # 3. generate image patch
    img_patch_cv, trans = generate_image_patch_cv2(cvimg,
                                                   center_x, center_y,
                                                   width, height,
                                                   patch_width, patch_height,
                                                   False, scale, rot,
                                                   border_mode=border_mode)
    if do_flip:
        flipped_img_patch_cv, flipped_trans = generate_image_patch_cv2(cvimg,
                                                                       center_x, center_y,
                                                                       width, height,
                                                                       patch_width, patch_height,
                                                                       True, scale, rot,
                                                                       border_mode=border_mode)

    image = img_patch_cv.copy()

    if is_bgr:
        image = image[:, :, ::-1]
    img_patch = convert_cvimg_to_tensor(image)
    img_patch_unnormalized = img_patch.copy()

    for n_c in range(min(img_channels, 3)):
        if mean is not None and std is not None:
            img_patch[n_c, :, :] = (
                img_patch[n_c, :, :] - mean[n_c]) / std[n_c]

    if do_flip:
        image_flipped = flipped_img_patch_cv.copy()
        if is_bgr:
            image_flipped = image_flipped[:, :, ::-1]
        flipped_img_patch_cv = image_flipped.copy()
        flipped_img_patch = convert_cvimg_to_tensor(image_flipped)
        flipped_img_patch_unnormalized = flipped_img_patch.copy()

        for n_c in range(min(img_channels, 3)):
            flipped_img_patch[n_c, :, :] = np.clip(
                flipped_img_patch[n_c, :, :] * color_scale[n_c], 0, 255)
            if mean is not None and std is not None:
                flipped_img_patch[n_c, :, :] = (
                    flipped_img_patch[n_c, :, :] - mean[n_c]) / std[n_c]

    keypoints_2d[:, :2] = trans_points2d_parallel(keypoints_2d[:, 0:2], trans)
    keypoints_2d[:, :-1] = keypoints_2d[:, :-1] / patch_width - 0.5

    if do_flip:
        return img_patch, flipped_img_patch, img_patch_unnormalized, flipped_img_patch_unnormalized, keypoints_2d, img_size, center_x, center_y, width, height, scale
    else:
        return img_patch, img_patch_unnormalized, keypoints_2d, img_size, center_x, center_y, width, height, scale


def get_examples(n_samples, mode, device, cam_model=None):
    """
    n_samples: number of samples to return
    mode: 'iid' or 'same_clip'
    """
    segmentor = YOLO('yolo11x-seg.pt')
    segmentor.to(device)

    data = np.load(_C.PATHS.DATA_EMDB_TESTLABELS, allow_pickle=True)
    imgnames = data['imgname']
    scales = data['scale']
    centers = data['center']
    poses = data['pose_cam'][:, :24*3].astype(float)
    betas = data['shape'].astype(float)[:, :10]
    keypoints_2d_all = data['gtkps'][:, :44]
    genders = data['gender']
    genders = np.array([0 if str(g) == 'm' or str(g)=='male'
                                    else 1 for g in genders]).astype(np.int32)
    lengths = scales.shape[0]

    normalize_img = Normalize(mean=_C.IMAGE_MEAN, std=_C.IMAGE_STD)
    
    n_video_frames = {}
    for imgname in imgnames:
        vid_name = f"{imgname.split(os.sep)[-4]}_{imgname.split(os.sep)[-3]}"
        if vid_name in n_video_frames:
            n_video_frames[vid_name] += 1
        else:
            n_video_frames[vid_name] = 1
            
    n_sample_per_video = n_samples // NUM_VIDEOS

    vid_names = {}
    samples, gts = [], []
    for i in tqdm(range(lengths)):
        scale = scales[i]
        center = centers[i]
        gender = genders[i]
        keypoints_2d = keypoints_2d_all[i]
        center_x, center_y = center
        bbox_size = expand_to_aspect_ratio(
            scale*200, target_aspect_ratio=None).max()
        assert bbox_size >= 1, f"bbox size is {bbox_size}"

        img_path = os.path.join(_C.PATHS.DATA_EMDB, imgnames[i])

        vid_name = f"{img_path.split(os.sep)[-4]}_{img_path.split(os.sep)[-3]}"

        if vid_name in vid_names:
            if mode == 'iid' and vid_names[vid_name] < n_sample_per_video:
                # Check if this frame should be sampled based on uniform spacing
                frame_id = int(img_path.split(os.sep)[-1].split('.')[0].replace("image_", ""))
                total_frames = n_video_frames[vid_name]
                step = total_frames // n_sample_per_video
                target_frames = [j * step for j in range(n_sample_per_video)]
                if frame_id not in target_frames:
                    continue
            elif mode == 'iid':
                continue
        else:
            vid_names[vid_name] = 0
        
        vid_names[vid_name] += 1
        frame_id = int(img_path.split(os.sep)
                       [-1].split('.')[0].replace("image_", ""))

        cv_img = cv2.imread(img_path, cv2.IMREAD_COLOR |
                            cv2.IMREAD_IGNORE_ORIENTATION)
        cv_img = cv_img[:, :, ::-1]
        aspect_ratio, img_full_resized = resize_image(cv_img,  _C.IMAGE_SIZE)
        img_full_resized = np.transpose(img_full_resized.astype('float32'),
                                        (2, 0, 1))/255.0
        img_full_resized = normalize_img(
            torch.from_numpy(img_full_resized).float()).to(device)

        img_patch_rgba, \
            flipped_img_patch_rgba, \
            img_patch_unnormalized, \
            flipped_img_patch_unnormalized, \
            keypoints_2d, \
            img_size, cx, cy, bbox_w, bbox_h, scale_aug = get_example(img_path,
                                                                      center_x, center_y,
                                                                      bbox_size, bbox_size,
                                                                      keypoints_2d,
                                                                      _C.IMAGE_SIZE, _C.IMAGE_SIZE,
                                                                      255. *
                                                                      np.array(
                                                                          _C.IMAGE_MEAN), 255. * np.array(_C.IMAGE_STD),
                                                                      is_bgr=True, border_mode=cv2.BORDER_CONSTANT,
                                                                      do_flip=True
                                                                      )

        segmentor_results = segmentor(torch.from_numpy(img_patch_unnormalized).unsqueeze(
            0).to(device) / 255,  verbose=False, conf=0.1)
        for idx in range(len(segmentor_results)):
            person_indices = (
                segmentor_results[idx].boxes.cls == PERSON_ID).nonzero()
            if len(person_indices) >= 1:
                person_mask = segmentor_results[idx].masks.data[person_indices].any(
                    dim=(0, 1)).float().to(device)
            else:
                person_mask = torch.zeros(
                    (img_patch_rgba.shape[1], img_patch_rgba.shape[2])).to(device)
            person_mask_downsampled = person_mask.unsqueeze(0)[..., 32:-32]
            person_mask_downsampled = F.interpolate(
                person_mask_downsampled.unsqueeze(1),
                size=(16, 12),
                mode='bilinear',
                align_corners=False
            ).squeeze(1)
            person_mask_downsampled = person_mask_downsampled.flatten(1)

        img_h, img_w = img_size[:2]
        bbox = torch.tensor([cx, cy, bbox_w * scale_aug]).float()
        res = torch.tensor([img_h, img_w]).float()
        img_patch = torch.from_numpy(
            img_patch_rgba[:3, :, :]).unsqueeze(0).to(device)
        samples.append({
            'img': img_patch,
            'person_masks': person_mask.unsqueeze(0),
            'person_masks_downsampled': person_mask_downsampled,
            'box_center': bbox[:2].unsqueeze(0).to(device),
            'box_size': bbox[2].unsqueeze(0).to(device),
            'img_size': res.unsqueeze(0).to(device),
        })
        
        gts.append({
            'gender': 'male' if gender == 0 else 'female',
            'smpl_params' : {'global_orient': torch.from_numpy(poses[i][:3]).unsqueeze(0).float().to(device),
                        'body_pose': torch.from_numpy(poses[i][3:]).unsqueeze(0).float().to(device),
                        'betas': torch.from_numpy(betas[i]).unsqueeze(0).float().to(device)
                        }
        })

        if cam_model is not None:
            cam, features = cam_model(img_full_resized.unsqueeze(0))
            vfov = cam[:, 1]
            fl_h = (img_h / (2 * torch.tan(vfov / 2))).item()
            cam_int = np.array(
                [[fl_h, 0, img_w/2.], [0, fl_h, img_h / 2.], [0, 0, 1]]).astype(np.float32)
            samples[-1]['cam_int'] = torch.from_numpy(
                cam_int).unsqueeze(0).to(device)

        if len(samples) == n_samples:
            break

    return samples, gts


def draw_cka_matrix(cka_matrix, model_name, example_type):
    model_label = {
        'camerahmr': 'CameraHMR',
        'hmr2': 'HMR2',
    }
    cka_matrix = cka_matrix.cpu().numpy()
    N = cka_matrix.shape[0]

    plt.figure(figsize=(8, 6))
    plt.imshow(cka_matrix, cmap='viridis', vmin=0, vmax=1)
    plt.colorbar()
    # plt.title(f"{model_label[model_name]} ({example_type} samples)")

    # Set tick spacing: adjust tick_spacing as needed
    tick_spacing = 2  
    ticks = list(range(0, N, tick_spacing))
    tick_labels = [i+1 for i in ticks]
    plt.xticks(ticks=ticks, labels=tick_labels)
    plt.yticks(ticks=ticks, labels=tick_labels)

    plt.tight_layout()
    plt.savefig(f"cka_{model_name}_{example_type}.png", dpi=300)
    plt.close()


def merge_layers_return_model(model, low_lay, high_lay, weight_factor):
    if low_lay < 0 or high_lay >= len(model.backbone.blocks):
        raise ValueError("The layer's index is out of range for the model")

    model_copy = deepcopy(model)

    for current_layer_idx in range(low_lay + 1, high_lay + 1):

        for projection in ['fc1', 'fc2']:
            model_copy.backbone.blocks[low_lay].mlp.__getattr__(projection).weight.data.add_(
                (model.backbone.blocks[current_layer_idx].mlp.__getattr__(projection).weight.data - model_copy.backbone.blocks[low_lay].mlp.__getattr__(projection).weight.data) * weight_factor
            )

        for projection in ['qkv', 'proj']:
            model_copy.backbone.blocks[low_lay].attn.__getattr__(projection).weight.data.add_(
                (model.backbone.blocks[current_layer_idx].attn.__getattr__(projection).weight.data - model_copy.backbone.blocks[low_lay].attn.__getattr__(projection).weight.data) * weight_factor
            )

    for current_layer_idx in range(high_lay, low_lay, -1):
        del(model_copy.backbone.blocks[current_layer_idx])


    return model_copy

def cal_sim(hidden_states_baseline, tmp_merged_extractor, examples):
    cka_vals = []
    for idx, example in enumerate(tqdm(examples)):
        tmp_merged_extractor(example)

        hidden_states1 = hidden_states_baseline[idx]

        hidden_states2 = tmp_merged_extractor.backbone.hidden_states
        masks = example['person_masks_downsampled']
        hidden_states2 = hidden_states2[masks.bool()].unsqueeze(0)

        cka = compute_cka(hidden_states1, hidden_states2)
        cka_vals.append(cka.item())
    sim_value = np.mean(cka_vals)
    return sim_value


def extract_mpjpe(extractor, examples, gts, body_models):
    mpjpe_list = []
    for example, gt in tqdm(zip(examples, gts)):
        out = extractor(example)
        out_smpl_params = out[0]
        out = body_models['neutral'](global_orient=out_smpl_params['global_orient'].reshape(1, 1, 3, 3),
                                     body_pose=out_smpl_params['body_pose'].reshape(
                                         1, -1, 3, 3),
                                     betas=out_smpl_params['betas'], pose2rot=False)
        out_pred_joints = out.joints[:, :24].reshape(-1, 3)
        
        out_gt = body_models[gt['gender']](**gt['smpl_params'])
        out_gt_joints = out_gt.joints[:, :24].reshape(-1, 3)
        
        mpjpe = torch.mean(torch.norm(out_pred_joints - out_gt_joints, dim=-1)).item()
        mpjpe_list.append(mpjpe)

    return np.mean(mpjpe_list) * 1000


@torch.no_grad()
def main():
    args = parse_args()

    device = torch.device(
        'cuda') if torch.cuda.is_available() else torch.device('cpu')

    print("[INFO] Initializing model...")
    extractor = init_model(args.model_name, args.type, device)
    if args.model_name == 'camerahmr':
        cam_model = init_cam_model(device)
    else:
        cam_model = None
    print("[INFO] Model initialized")

    print("[INFO] Getting examples...")
    examples, gts = get_examples(
        args.n_samples, args.example_type, device, cam_model)
    print(f"[INFO] {len(examples)} Examples obtained")

    if args.type == 'analysis':
        cka_matrix = torch.zeros((NUM_LAYERS, NUM_LAYERS), device=device)
        print("[INFO] Computing CKA matrix...")
        for example in tqdm(examples):
            extractor(example)
            hidden_states = extractor.backbone.hidden_states
            layer_indices = list(hidden_states.keys())
            for layer_idx1 in layer_indices:
                for layer_idx2 in layer_indices:
                    if layer_idx1 < layer_idx2:
                        cka = compute_cka(hidden_states[layer_idx1],
                                          hidden_states[layer_idx2])
                        cka_matrix[layer_idx1, layer_idx2] += cka
                        cka_matrix[layer_idx2, layer_idx1] += cka
                    elif layer_idx1 == layer_idx2:
                        cka = compute_cka(hidden_states[layer_idx1],
                                          hidden_states[layer_idx2])
                        cka_matrix[layer_idx1, layer_idx2] += cka
        cka_matrix /= args.n_samples
        print("[INFO] CKA matrix computed")

        draw_cka_matrix(cka_matrix, args.model_name, args.example_type)
    elif args.type == 'eclm':
        genders = ['male', 'female', 'neutral']
        body_models = {
            gender: SMPL(model_path=_C.BMODEL.SMPL, gender=gender).to(device)
            for gender in genders
        }
        for body_model in body_models.values():
            body_model.eval()
            body_model.requires_grad_(False)

        high_lay = NUM_LAYERS - 1
        low_lay = high_lay - 1
        mpjpe_baseline = extract_mpjpe(
            extractor, examples, gts, body_models)
        print(f"[INFO] MPJPE baseline: {mpjpe_baseline:.4f}")
        while low_lay >= 0:
            tmp_merged_extractor = merge_layers_return_model(
                extractor, low_lay, high_lay, weight_factor=1)

            mpjpe_merged = extract_mpjpe(
                tmp_merged_extractor, examples, gts, body_models)
            print(
                f"[INFO] MPJPE after merging layers {low_lay} to {high_lay}: {mpjpe_merged:.4f}")
            if (mpjpe_merged - mpjpe_baseline) < args.threshold:
                low_lay -= 1
            else:
                if low_lay + 1 != high_lay:
                    extractor = merge_layers_return_model(
                        extractor, low_lay + 1, high_lay, weight_factor=1)

                    print(f"[INFO] Merged layers {low_lay + 1} to {high_lay}")
                    high_lay = low_lay
                    low_lay = high_lay - 1
                else:
                    high_lay -= 1
                    low_lay -= 1
        print("[INFO] number of layers after merging:",
              len(extractor.backbone.blocks))
        torch.save({
            'n_blocks': len(extractor.backbone.blocks),
            'weights': extractor.state_dict(),
        }, f"checkpoint/merged_{args.model_name}_thresh{args.threshold}_{args.n_samples}.pth")
    else:
        raise ValueError(f"Type {args.type} not supported.")


if __name__ == "__main__":
    main()