#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import List
import cv2
import numpy
__author__ = "Christian Heider Nielsen"
__doc__ = r"""
Created on 10/10/2019
"""
__all__ = ["draw_cube", "draw_axis", "cube_3d_matrix"]
from draugr.opencv_utilities.windows.image import show_image
from warg import Number
[docs]def draw_axis(
img: numpy.ndarray,
corners: numpy.ndarray,
rotation_vectors,
translation_vectors,
camera_matrix,
dist_coef,
size: Number = 1,
) -> numpy.ndarray:
"""
:param img:
:type img:
:param corners:
:type corners:
:param rotation_vectors:
:type rotation_vectors:
:param translation_vectors:
:type translation_vectors:
:param camera_matrix:
:type camera_matrix:
:param dist_coef:
:type dist_coef:
:param size:
:type size:
:return:
:rtype:"""
axis_size: numpy.ndarray = numpy.eye(3, 3) @ numpy.diag([size, size, -size])
axis_size = axis_size.astype(numpy.float32)
img_pts, jac = cv2.projectPoints(
axis_size, rotation_vectors, translation_vectors, camera_matrix, dist_coef
) # project 3D points to image plane
corner = tuple(corners[0].ravel())
img = cv2.line(img, corner, tuple(img_pts[0].ravel()), (255, 0, 0), 5)
img = cv2.line(img, corner, tuple(img_pts[1].ravel()), (0, 255, 0), 5)
img = cv2.line(img, corner, tuple(img_pts[2].ravel()), (0, 0, 255), 5)
return img
[docs]def cube_3d_matrix() -> List[List[int]]:
"""
:return:
:rtype:"""
return [
[0, 0, 0],
[0, 1, 0],
[1, 1, 0],
[1, 0, 0],
[0, 0, -1],
[0, 1, -1],
[1, 1, -1],
[1, 0, -1],
]
[docs]def draw_cube(
img: numpy.ndarray,
rotation_vectors: numpy.ndarray = None,
translation_vectors: numpy.ndarray = None,
camera_matrix: numpy.ndarray = None,
dist_coef: float = None,
size: Number = 6,
) -> numpy.ndarray:
"""
:param img:
:type img:
:param rotation_vectors:
:type rotation_vectors:
:param translation_vectors:
:type translation_vectors:
:param camera_matrix:
:type camera_matrix:
:param dist_coef:
:type dist_coef:
:param size:
:type size:
:return:
:rtype:"""
if rotation_vectors is None:
rotation_vectors = numpy.eye(3, 3)
if translation_vectors is None:
translation_vectors = numpy.zeros(3)
if camera_matrix is None:
camera_matrix = numpy.eye(3, 3)
if dist_coef is None:
dist_coef = numpy.zeros(5)
cube_size = numpy.float32(cube_3d_matrix()) * size
img_pts, jac2 = cv2.projectPoints(
cube_size, rotation_vectors, translation_vectors, camera_matrix, dist_coef
)
img_pts = numpy.int32(img_pts).reshape(-1, 2)
# draw ground floor in green
img = cv2.drawContours(img, [img_pts[:4]], -1, (0, 255, 0), -3)
# draw pillars in blue color
for i, j in zip(range(4), range(4, 8)):
img = cv2.line(img, tuple(img_pts[i]), tuple(img_pts[j]), (255, 0, 0), 3)
# draw top layer in red color
img = cv2.drawContours(img, [img_pts[4:]], -1, (0, 0, 255), 3)
return img
if __name__ == "__main__":
show_image(
draw_cube(
numpy.zeros((50, 50)),
),
wait=True,
)