detector.py 1.55 KB
Newer Older
Tom Pillot's avatar
Tom Pillot committed
1
2
3
4
5
import apriltag
import cv2


class Detector(apriltag.Detector):
6
7
    cam = cv2.VideoCapture(0)

Tom Pillot's avatar
Tom Pillot committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
    def __init__(self):
        super().__init__()
        self.DECISION_MARGIN = 40

    def detect(self, img, return_image=False):
        dets = super().detect(img)
        for det in dets:
            if det.decision_margin > self.DECISION_MARGIN:
                return det.tag_id

    def detect_from_cam(self):
        cv2.namedWindow("Detection")

        tag_id = None
        while not tag_id:
23
            ret, frame = self.cam.read()
Tom Pillot's avatar
Tom Pillot committed
24
            if not ret:
25
                print("Failed to grab frame.")
Tom Pillot's avatar
Tom Pillot committed
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
                break

            cv2.imshow("Detection", frame)

            k = cv2.waitKey(1)
            if k % 256 == 27:
                # ESC pressed
                print("Escape hit, closing...")
                break

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            tag_id = self.detect(gray)

        cv2.destroyAllWindows()
        return tag_id
41
42
43
44
45
46
47
48
49
50
51
52
53
54

    def show_tags_on_image(self, img):
        greys = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        dets = super().detect(greys)
        for det in dets:
            if det.decision_margin >= self.DECISION_MARGIN:
                rect = det.corners.astype(int).reshape((-1, 1, 2))
                cv2.polylines(img, [rect], True, (0, 0, 255), 2)
                ident = str(det.tag_id)
                pos = det.center.astype(int) + (-10, 10)
                cv2.putText(img, ident, tuple(pos), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    def stop_webcam(self):
        self.cam.release()