movement_bc_test

import numpy as np
import cv2
import math
import json
import re

path_to_files = '/Users/shtern/runner-new-motion'
movement_json_path = '/Users/shtern/Downloads/movement.json'
res_prod_map = {2773: 'd2mj5veol0', 3139: '493lkx9o7e', 3113: 'y9m6nnl3e2', 2938: 'vjmkdv93qy'}
extra_json_name = 'scan_0_extra.json'
marker_size_cm = 2.13
num_points = 32


class MotionData(object):
    def __init__(self, roll, pitch, yaw):
        self.roll = roll
        self.pitch = pitch
        self.yaw = yaw
        self.rotation_matrix = self.eulers_to_rotation_matrix()

    @classmethod
    def from_dict(cls, d):
        return cls(d["attitude"]["roll"], d["attitude"]["pitch"], d["attitude"]["yaw"])

    def eulers_to_rotation_matrix(self):
        m11 = np.cos(self.pitch) * np.cos(self.yaw)
        m21 = np.cos(self.pitch) * np.sin(self.yaw)
        m31 = -np.sin(self.pitch)

        m12 = np.sin(self.roll) * np.sin(self.pitch) * np.cos(self.yaw) - np.cos(self.roll) * np.sin(self.yaw)
        m22 = np.cos(self.roll) * np.cos(self.yaw) + np.sin(self.roll) * np.sin(self.pitch) * np.sin(self.yaw)
        m32 = np.sin(self.roll) * np.cos(self.pitch)

        m13 = np.sin(self.roll) * np.sin(self.yaw) + np.cos(self.roll) * np.sin(self.pitch) * np.cos(self.yaw)
        m23 = np.cos(self.roll) * np.sin(self.pitch) * np.cos(self.yaw) - np.sin(self.roll) * np.cos(self.yaw)
        m33 = np.cos(self.roll) * np.cos(self.pitch)
        return np.array([m11, m12, m13, m21, m22, m23, m31, m32, m33]).reshape(3, 3)


class Center(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y

    @classmethod
    def from_list(cls, lst):
        return cls(lst[0], lst[1])

    def to_list(self):
        return [self.x, self.y]

    @classmethod
    def from_dict(cls, d):
        return cls(d["x"], d["y"])

    def to_dict(self):
        return {"x": self.x, "y": self.y}

    def __str__(self):
        return str(self.to_dict())


class Size(object):
    def __init__(self, width, height):

        self.width = width
        self.height = height

    @classmethod
    def from_list(cls, lst):
        return cls(lst[0], lst[1])

    @classmethod
    def from_dict(cls, d):
        return cls(d["width"], d["height"])

    def to_dict(self):
        return {"width": self.width, "height": self.height}

    def __str__(self):
        return str(self.to_dict())


class Marker(object):
    def __init__(self, center, size, angle, sharpness):

        self.center = center
        self.size = size
        self.angle = angle  # The rotation angle in a clockwise direction
        self.sharpness = sharpness

    @classmethod
    def from_list(cls, lst):
        return cls(Center.from_list(lst[:2]), Size.from_list(lst[2:4]), lst[4], lst[5])

    def to_list(self):
        return [
            self.center.x,
            self.center.y,
            self.size.width,
            self.size.height,
            self.angle,
            self.sharpness,
        ]

    @classmethod
    def from_dict(cls, d):
        return cls(
            Center.from_dict(d["center"]),
            Size.from_dict(d["size"]),
            d["angle"],
            d["sharpness"],
        )

    def to_dict(self):
        return {
            "center": self.center.to_dict(),
            "size": self.size.to_dict(),
            "angle": self.angle,
            "sharpness": self.sharpness,
        }

    def __str__(self):
        return str(self.to_dict())

    def __repr__(self):
        return str(self)

    def axis_size(self):
        return max([self.size.width, self.size.height]) / 2

    def draw(self, img, filled_white=False):
        center = (int(self.center.x), int(self.center.y))
        dimensions = (int(self.size.width * 0.63), int(self.size.height * 0.63))
        if filled_white:
            color = (255, 255, 255)
            img = cv2.ellipse(img, center, dimensions, self.angle, 0, 360, color, -1)
        else:
            color = (255, 128, 0)
            img = cv2.ellipse(img, center, dimensions, self.angle, 0, 360, color, 4)
        return img

    def find_major_axis_px(self):

        # marker.angle in a clockwise direction.
        # Add 90 degrees to match the sin/cos calculation of the marker major axis points.
        angle = np.deg2rad(self.angle + 90)

        sin_angle, cos_angle = np.sin(angle), np.cos(angle)
        axis_size = self.axis_size()

        a_axis_px_1 = [
            int(self.center.x + axis_size * cos_angle),
            int(self.center.y + axis_size * sin_angle),
        ]
        a_axis_px_2 = [
            int(self.center.x - axis_size * cos_angle),
            int(self.center.y - axis_size * sin_angle),
        ]
        x_diff = int(self.center.x - axis_size * cos_angle) - int(self.center.x + axis_size * cos_angle)
        y_diff = int(self.center.y - axis_size * sin_angle) - int(self.center.y + axis_size * sin_angle)
        return math.hypot(x_diff, y_diff)


class MotionSolution:
    def __init__(self, f, C, motion_data):
        self.scan_plain_markers = []
        self.scan_markers_points = []
        self.scan_3d_markers_points = []
        self.f = f
        self.C = C
        self.motion_data = motion_data
        self.a_mat = np.zeros((2, 2))
        self.b_mat = np.zeros((1, 2))

    def append_markers_data(self, markers, marker_points, marker_points_3d):
        self.scan_plain_markers.append(markers)
        self.scan_markers_points.append([point*2 for point in marker_points])
        self.scan_3d_markers_points.append(marker_points_3d)

    def solve(self):
        temp = np.linalg.pinv(self.a_mat.transpose().dot(self.a_mat)).dot(self.a_mat.transpose())
        return temp.dot(self.b_mat)

    def build_matrices(self):
        num_records = len(self.scan_plain_markers)

        self.a_mat = np.zeros((num_records*2*num_points, 3*num_points+num_records*2))
        self.b_mat = np.zeros((num_records*2*num_points, 1))
        current_line_index = -1
        for frame_num in range(len(self.scan_plain_markers)):
            rotation_matrix = self.motion_data[frame_num].rotation_matrix
            markers_points_frame = self.scan_markers_points[frame_num]
            plain_markers = self.scan_plain_markers[frame_num]
            if len(plain_markers) == 0:
                continue

            for marker_index, (marker, marker_points) in enumerate(zip(plain_markers, markers_points_frame)):
                scale = get_scale(marker)
                for point_index, marker_point in enumerate(marker_points):
                    for plain_coord_index, marker_coordinate in enumerate(marker_point):
                        current_line_index += 1
                        for m in range(3):
                            self.a_mat[current_line_index, marker_index * 3 * len(marker_points) + point_index*3 + m] =\
                                scale * self.f * rotation_matrix[plain_coord_index, m]
                        self.a_mat[current_line_index, num_points*3 + 2*frame_num + plain_coord_index] = -scale * self.f
                        self.b_mat[current_line_index, 0] = marker_coordinate + scale * self.C[marker_index]


def get_scale(marker):
    return marker_size_cm / marker.find_major_axis_px()


def get_motion_data(path):
    with open(path) as f:
        extra_json = json.load(f)
        return extra_json['motion']

def test():
    with open(movement_json_path) as f:
        movement_json_data = json.load(f)

    for res_id, data in movement_json_data.items():
        principal_point = data['principal_point']
        px_focal_length = data['px_focal_length']
        camera_matrix = np.array(data['camera_matrix'])
        motion_data = get_motion_data(path_to_files+f'/{res_prod_map[int(res_id)]}/{extra_json_name}')
        solution = MotionSolution(px_focal_length, principal_point)
        r = re.compile('frame[0-9]*')

        new_list = list(filter(r.match, data.keys()))
        new_list.sort()
        numbers = [int(x[5:]) for x in new_list]
        for num in numbers:

            rec_data = data[f'frame{str(num).zfill(4)}']
            if not rec_data:
                break
            else:
                solution.append_data(np.array(rec_data['R']), [Marker.from_dict(x) for x in rec_data['markers_2d']],
                                     np.array(rec_data['markers_2d_16_points']), np.array(rec_data['markers_3d_16_points']))

        solution.build_matrices()
        print(solution.solve())
        print(solution.scan_3d_markers_points)


if __name__ == '__main__':
    test()