clean/video/fakestormer/package_utils/utils.py

# -*- coding: utf-8 -*-
import logging
import os
from copy import deepcopy

import cv2
import numpy as np
import plotly.graph_objects as go
import simplejson as json
import torch
from losses.losses import _sigmoid


def file_extention(file_path):
    f_name, f_extension = os.path.splitext(file_path)
    return f_name, f_extension


def make_dir(dir_path):
    if not os.path.exists(dir_path):
        os.mkdir(dir_path)


def vis_heatmap(images, heatmaps, file_name, **kwargs):
    temp_locs = kwargs.get("temp_loc_preds")

    # hm_h, hm_w = heatmaps.shape[1:]
    hm_h, hm_w = np.array(images[0]).shape[:2]

    masked_image = np.zeros((hm_h, hm_w * heatmaps.shape[0], 3), dtype=np.uint8)

    for i in range(heatmaps.shape[0]):
        heatmap = heatmaps[i]
        heatmap = np.clip(heatmap * 255, 0, 255).astype(np.uint8)
        heatmap = np.squeeze(heatmap)

        # heatmap_h = heatmap.shape[0]
        # heatmap_w = heatmap.shape[1]

        if isinstance(images, list):
            # resized_image = cv2.resize(np.array(images[i]), (int(heatmap_h), int(heatmap_w)))
            heatmap = cv2.resize(
                heatmap, np.array(images[i]).shape[:2], interpolation=cv2.INTER_LINEAR
            )
            colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
            masked_image[:, hm_w * i : hm_w * (i + 1), :] = (
                colored_heatmap * 0.7 + np.array(images[i]) * 0.3
            )
        else:
            # resized_image = cv2.resize(images, (int(heatmap_h), int(heatmap_w)))
            heatmap = cv2.resize(
                heatmap, images.shape[:2], interpolation=cv2.INTER_LINEAR
            )
            colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
            masked_image[:, hm_w * i : hm_w * (i + 1), :] = (
                colored_heatmap * 0.7 + images * 0.3
            )

        if temp_locs is not None:
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 1
            color = (255, 255, 255)  # White color in BGR
            thickness = 2
            position = (20, 30)

            temp_loc = temp_locs[i]
            text = f"{temp_loc:.5f}"
            cv2.putText(
                masked_image[:, hm_w * i : hm_w * (i + 1), :],
                text,
                position,
                font,
                font_scale,
                color,
                thickness,
            )
    cv2.imwrite(file_name, masked_image)


def vis_3d_heatmap(heatmap, file_name):
    # Define the dimensions of the cuboid
    z_dim, y_dim, x_dim = heatmap.shape

    Z, Y, X = np.mgrid[:z_dim, :y_dim, :x_dim]

    fig = go.Figure(
        data=go.Volume(
            x=X.flatten(),
            y=Y.flatten(),
            z=Z.flatten(),
            value=heatmap.flatten(),
            isomin=0.0,
            isomax=0.999,
            opacity=0.1,
            surface_count=25,
        )
    )
    fig.update_layout(
        scene_xaxis_showticklabels=False,
        scene_yaxis_showticklabels=False,
        scene_zaxis_showticklabels=False,
    )
    fig.write_image(file_name)


def save_batch_heatmaps(
    batch_image, batch_heatmaps, file_name, normalize=True, batch_cls=None
):
    """
    batch_image: [batch_size, channel, height, width]
    batch_heatmaps: ['batch_size, num_joints, height, width]
    batch_cls: ['batch_size, num_joints, 1]
    file_name: saved file name
    """
    if normalize:
        batch_image = batch_image.clone()
        min = float(batch_image.min())
        max = float(batch_image.max())

        batch_image.add_(-min).div_(max - min + 1e-5)

    batch_size = batch_heatmaps.size(0)
    num_joints = batch_heatmaps.size(1)
    heatmap_height = batch_heatmaps.size(2)
    heatmap_width = batch_heatmaps.size(3)

    grid_image = np.zeros(
        (batch_size * heatmap_height, (num_joints + 1) * heatmap_width, 3),
        dtype=np.uint8,
    )

    for i in range(batch_size):
        if batch_image.dim() == 4:
            image = (
                batch_image[i]
                .mul(255)
                .clamp(0, 255)
                .byte()
                .permute(1, 2, 0)
                .cpu()
                .numpy()
            )
        else:
            image = (
                batch_image[i]
                .mul(255)
                .clamp(0, 255)
                .byte()
                .permute(1, 2, 3, 0)
                .cpu()
                .numpy()
            )
        heatmaps = batch_heatmaps[i].mul(255).clamp(0, 255).byte().cpu().numpy()

        height_begin = heatmap_height * i
        height_end = heatmap_height * (i + 1)
        for j in range(num_joints):
            if image.ndim == 4:
                image = image[j, :, :, :]

            resized_image = cv2.resize(image, (int(heatmap_width), int(heatmap_height)))

            heatmap = heatmaps[j, :, :]
            colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)

            if batch_cls is not None:
                cls = batch_cls[i][j].detach().cpu().numpy()
                colored_heatmap = cv2.putText(
                    colored_heatmap,
                    f"Cls Pred: {cls}",
                    (heatmap_width * (j + 1) - 15, 10),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1,
                    1,
                    cv2.LINE_AA,
                )
            masked_image = colored_heatmap * 0.7 + resized_image * 0.3

            width_begin = heatmap_width * (j + 1)
            width_end = heatmap_width * (j + 2)

            grid_image[height_begin:height_end, width_begin:width_end, :] = masked_image

        grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image
    cv2.imwrite(file_name, grid_image)


def debugging_panel(
    debug_cfg,
    batch_image,
    batch_heatmaps_gt,
    batch_heatmaps_pred,
    idx,
    normalize=True,
    batch_cls_gt=None,
    batch_cls_pred=None,
    split="train",
):
    if debug_cfg.save_hm_gt:
        save_batch_heatmaps(
            batch_image,
            batch_heatmaps_gt,
            f"samples/{split}_debugs/hm_gt_{idx}.jpg",
            normalize=normalize,
        )

    if debug_cfg.save_hm_pred:
        batch_heatmaps_pred_ = _sigmoid(batch_heatmaps_pred.clone())
        save_batch_heatmaps(
            batch_image,
            batch_heatmaps_pred_,
            f"samples/{split}_debugs/hm_pred_{idx}.jpg",
            normalize=normalize,
        )


def save_file(data, file_path):
    f_name, f_extention = file_extention(file_path)

    if f_extention == ".json":
        with open(file_path, "w") as f:
            json.dump(data, f)
        print(f"Data has been saved to --- {file_path}")
    else:
        raise ValueError(f"{f_extention} is not supported now!")


def load_file(file_path):
    f_name, f_extention = file_extention(file_path)

    if f_extention == ".json":
        with open(file_path, "r") as f:
            data = json.load(f)
        print(f"Data has been loaded from --- {file_path}")
    else:
        raise ValueError(f"{f_extention} is not supported now!")

    return data


def draw_landmarks(image, landmarks):
    """This function is to draw facial landmarks into transformed images"""
    assert landmarks is not None, "Landmarks can not be None!"

    img_cp = deepcopy(image)
    landmarks = landmarks.astype(int)

    for i, p in enumerate(landmarks):
        img_cp = cv2.circle(img_cp, (p[0], p[1]), 2, (0, 255, 0), 1)

    return img_cp


def draw_most_vul_points(blended_mask):
    """Detecting and Drawing the most vulnerable points for visualization purpose"""
    b_mask_cp = deepcopy(blended_mask)
    target_H, target_W, target_C = b_mask_cp.shape

    max_val = b_mask_cp[..., 0].max()
    max_val = max_val if max_val > 0 else 1

    m_v_indices = np.where(b_mask_cp == max_val)

    for j, i in zip(m_v_indices[0], m_v_indices[1]):
        b_mask_cp[j, i] = (255, 0, 0)

    return b_mask_cp