aboutsummaryrefslogtreecommitdiff
path: root/src/argaze/ArUcoMarkers/ArUcoCube.py
blob: 88163d86880acc79e49f9de9c65842445dbff2e4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#!/usr/bin/env python

from typing import Tuple
from dataclasses import dataclass, field
import json
import math
import itertools

from argaze.ArUcoMarkers import ArUcoSet

import numpy
import cv2 as cv
import cv2.aruco as aruco

@dataclass
class ArUcoCube(ArUcoSet.ArUcoSet):
    """Define ArUco cube as a specific ArUco set."""

    edge_size: int = field(init=False)
    """Size of the cube edges in centimeter."""

    def __init__(self, configuration_filepath):
        """Define cube from a .json  file."""

        # Load generic set configuration data
        super().__init__(configuration_filepath)

        # Load specific cube configuration data
        with open(configuration_filepath) as configuration_file:

            # Deserialize .json
            # TODO find a better way
            configuration = json.load(configuration_file)

            # Load edge size
            self.edge_size = configuration['edge_size']

    def draw(self, frame, K, D, draw_places=True):
        """Draw cube, axis and places."""

        l = self.edge_size / 2
        ll = self.edge_size

        # Select color according validity score
        n = 95 * self._validity if self._validity < 2 else 0
        f = 159 * self._validity if self._validity < 2 else 255

        # Draw left face
        leftPoints = numpy.float32([[-l, l, l], [-l, -l, l], [-l, -l, -l], [-l, l, -l]]).reshape(-1, 3)
        leftPoints, _ = cv.projectPoints(leftPoints, self._rotation, self._translation, K, D)
        leftPoints = leftPoints.astype(int)
        
        cv.line(frame, tuple(leftPoints[0].ravel()), tuple(leftPoints[1].ravel()), (n,n,f), 2)
        cv.line(frame, tuple(leftPoints[1].ravel()), tuple(leftPoints[2].ravel()), (n,n,f), 2)
        cv.line(frame, tuple(leftPoints[2].ravel()), tuple(leftPoints[3].ravel()), (n,n,f), 2)
        cv.line(frame, tuple(leftPoints[3].ravel()), tuple(leftPoints[0].ravel()), (n,n,f), 2)

        # Draw top face
        topPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, l, -l], [l, l, -l]]).reshape(-1, 3)
        topPoints, _ = cv.projectPoints(topPoints, self._rotation, self._translation, K, D)
        topPoints = topPoints.astype(int)

        cv.line(frame, tuple(topPoints[0].ravel()), tuple(topPoints[1].ravel()), (n,f,n), 2)
        cv.line(frame, tuple(topPoints[1].ravel()), tuple(topPoints[2].ravel()), (n,f,n), 2)
        cv.line(frame, tuple(topPoints[2].ravel()), tuple(topPoints[3].ravel()), (n,f,n), 2)
        cv.line(frame, tuple(topPoints[3].ravel()), tuple(topPoints[0].ravel()), (n,f,n), 2)

        # Draw front face
        frontPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, -l, l], [l, -l, l]]).reshape(-1, 3)
        frontPoints, _ = cv.projectPoints(frontPoints, self._rotation, self._translation, K, D)
        frontPoints = frontPoints.astype(int)
        
        cv.line(frame, tuple(frontPoints[0].ravel()), tuple(frontPoints[1].ravel()), (f,n,n), 2)
        cv.line(frame, tuple(frontPoints[1].ravel()), tuple(frontPoints[2].ravel()), (f,n,n), 2)
        cv.line(frame, tuple(frontPoints[2].ravel()), tuple(frontPoints[3].ravel()), (f,n,n), 2)
        cv.line(frame, tuple(frontPoints[3].ravel()), tuple(frontPoints[0].ravel()), (f,n,n), 2)

        # Draw axis and places
        super().draw(frame, K, D, draw_places)