utils.py

import numpy as np
import torch

# from torch.utils.serialization import load_lua
import os

import scipy.io as sio
import cv2
from math import cos, sin


def softmax_temperature(tensor, temperature):
    result = torch.exp(tensor / temperature)
    result = torch.div(result, torch.sum(result, 1).unsqueeze(1).expand_as(result))
    return result


def get_pose_params_from_mat(mat_path):
    # This functions gets the pose parameters from the .mat
    # Annotations that come with the Pose_300W_LP dataset.
    mat = sio.loadmat(mat_path)
    # [pitch yaw roll tdx tdy tdz scale_factor]
    pre_pose_params = mat["Pose_Para"][0]
    # Get [pitch, yaw, roll, tdx, tdy]
    pose_params = pre_pose_params[:5]
    return pose_params


def get_ypr_from_mat(mat_path):
    # Get yaw, pitch, roll from .mat annotation.
    # They are in radians
    mat = sio.loadmat(mat_path)
    # [pitch yaw roll tdx tdy tdz scale_factor]
    pre_pose_params = mat["Pose_Para"][0]
    # Get [pitch, yaw, roll]
    pose_params = pre_pose_params[:3]
    return pose_params


def get_pt2d_from_mat(mat_path):
    # Get 2D landmarks
    mat = sio.loadmat(mat_path)
    pt2d = mat["pt2d"]
    return pt2d


def mse_loss(input, target):
    return torch.sum(torch.abs(input.data - target.data) ** 2)


def plot_pose_cube(img, yaw, pitch, roll, tdx=None, tdy=None, size=150.0):
    # Input is a cv2 image
    # pose_params: (pitch, yaw, roll, tdx, tdy)
    # Where (tdx, tdy) is the translation of the face.
    # For pose we have [pitch yaw roll tdx tdy tdz scale_factor]

    p = pitch * np.pi / 180
    y = -(yaw * np.pi / 180)
    r = roll * np.pi / 180
    if tdx is not None and tdy is not None:
        face_x = tdx - 0.50 * size
        face_y = tdy - 0.50 * size
    else:
        height, width = img.shape[:2]
        face_x = width / 2 - 0.5 * size
        face_y = height / 2 - 0.5 * size

    x1 = size * (cos(y) * cos(r)) + face_x
    y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y
    x2 = size * (-cos(y) * sin(r)) + face_x
    y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y
    x3 = size * (sin(y)) + face_x
    y3 = size * (-cos(y) * sin(p)) + face_y

    # Draw base in red
    cv2.line(img, (int(face_x), int(face_y)), (int(x1), int(y1)), (0, 0, 255), 3)
    cv2.line(img, (int(face_x), int(face_y)), (int(x2), int(y2)), (0, 0, 255), 3)
    cv2.line(
        img,
        (int(x2), int(y2)),
        (int(x2 + x1 - face_x), int(y2 + y1 - face_y)),
        (0, 0, 255),
        3,
    )
    cv2.line(
        img,
        (int(x1), int(y1)),
        (int(x1 + x2 - face_x), int(y1 + y2 - face_y)),
        (0, 0, 255),
        3,
    )
    # Draw pillars in blue
    cv2.line(img, (int(face_x), int(face_y)), (int(x3), int(y3)), (255, 0, 0), 2)
    cv2.line(
        img,
        (int(x1), int(y1)),
        (int(x1 + x3 - face_x), int(y1 + y3 - face_y)),
        (255, 0, 0),
        2,
    )
    cv2.line(
        img,
        (int(x2), int(y2)),
        (int(x2 + x3 - face_x), int(y2 + y3 - face_y)),
        (255, 0, 0),
        2,
    )
    cv2.line(
        img,
        (int(x2 + x1 - face_x), int(y2 + y1 - face_y)),
        (int(x3 + x1 + x2 - 2 * face_x), int(y3 + y2 + y1 - 2 * face_y)),
        (255, 0, 0),
        2,
    )
    # Draw top in green
    cv2.line(
        img,
        (int(x3 + x1 - face_x), int(y3 + y1 - face_y)),
        (int(x3 + x1 + x2 - 2 * face_x), int(y3 + y2 + y1 - 2 * face_y)),
        (0, 255, 0),
        2,
    )
    cv2.line(
        img,
        (int(x2 + x3 - face_x), int(y2 + y3 - face_y)),
        (int(x3 + x1 + x2 - 2 * face_x), int(y3 + y2 + y1 - 2 * face_y)),
        (0, 255, 0),
        2,
    )
    cv2.line(
        img,
        (int(x3), int(y3)),
        (int(x3 + x1 - face_x), int(y3 + y1 - face_y)),
        (0, 255, 0),
        2,
    )
    cv2.line(
        img,
        (int(x3), int(y3)),
        (int(x3 + x2 - face_x), int(y3 + y2 - face_y)),
        (0, 255, 0),
        2,
    )

    return img


def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size=100):

    pitch = pitch * np.pi / 180
    yaw = -(yaw * np.pi / 180)
    roll = roll * np.pi / 180

    if tdx is not None and tdy is not None:
        tdx = tdx
        tdy = tdy
    else:
        height, width = img.shape[:2]
        tdx = width / 2
        tdy = height / 2

    # X-Axis pointing to right. drawn in red
    x1 = size * (cos(yaw) * cos(roll)) + tdx
    y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy

    # Y-Axis | drawn in green
    #        v
    x2 = size * (-cos(yaw) * sin(roll)) + tdx
    y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy

    # Z-Axis (out of the screen) drawn in blue
    x3 = size * (sin(yaw)) + tdx
    y3 = size * (-cos(yaw) * sin(pitch)) + tdy

    cv2.line(img, (int(tdx), int(tdy)), (int(x1), int(y1)), (0, 0, 255), 3)
    cv2.line(img, (int(tdx), int(tdy)), (int(x2), int(y2)), (0, 255, 0), 3)
    cv2.line(img, (int(tdx), int(tdy)), (int(x3), int(y3)), (255, 0, 0), 2)

    return img


# helper functions to create filename.txt files
def are_angles_compatible(mat_file_path):
    import utils
    import numpy as np

    # Code copied from dataset.py

    # We get the pose in radians
    pose = utils.get_ypr_from_mat(mat_file_path)
    # And convert to degrees.
    pitch = pose[0] * 180 / np.pi
    yaw = pose[1] * 180 / np.pi
    roll = pose[2] * 180 / np.pi
    return (
        pitch < 99
        and pitch > -99
        and yaw < 99
        and yaw > -99
        and roll < 99
        and roll > -99
    )


def generate_filenames(dataset_path):
    import time

    found = 0
    using = 0
    start = time.time()
    with open(dataset_path + "_filenames.txt", "w") as textfile:
        for path, directories, files in os.walk(dataset_path):
            print(path)
            if len(files) > 0:  # folder with images
                # FIXME: dirty workaround
                subdirs = path.split("/")
                if "300" in dataset_path:
                    subdir = os.path.join(subdirs[-2], subdirs[-1])
                else:
                    subdir = subdirs[-1]
                counter = 0
                for imname in files:
                    if imname.endswith(".jpg"):
                        found += 1
                        if are_angles_compatible(
                            os.path.join(path, imname[:-4]) + ".mat"
                        ):
                            using += 1
                            textfile.write(os.path.join(subdir, imname[:-4]) + "\n")
    elapsed = time.time() - start
    print(
        "{}: Using/Found {}/{} ({:.1f} sec.)".format(
            dataset_path, using, found, elapsed
        )
    )