conan/processing/process_AprilTag.py
2021-10-17 21:32:30 +02:00

95 lines
2.7 KiB
Python

import os
from sys import platform
if platform == "linux" or platform == "linux2":
# linux
import apriltag
elif platform == "darwin":
# OS X
import apriltag
elif platform == "win32":
# Windows
import pupil_apriltags as apriltag
import cv2
import matplotlib.pyplot as plt
import pandas as pd
visualize = True
def process(file):
VIDEO = file
VIDEOOUT = VIDEO.split("/")[-1].split(".")[0]
ROOT = "/".join(VIDEO.split("/")[:-1]) + "/"
TMP_DIR = "/".join(VIDEO.split("/")[:-2]) + "/temp/"
FRAMES = "%s%s_frames" % (TMP_DIR, VIDEOOUT)
if not os.path.exists(FRAMES):
print('WARNING: Could not find frame directory')
return
img_paths = [f for f in os.listdir(FRAMES) if 'jpg' in f]
print('Number of frames: ', len(img_paths))
if platform == "linux" or platform == "linux2" or platform == "darwin":
# Circumvent error: too many borders in contour_detect (max of 32767!)
options = apriltag.DetectorOptions(refine_edges=False, quad_contours=False)
detector = apriltag.Detector(options)
elif platform == "win32":
print('WARNING: apriltag2 not supported on windows, running with pupil_apriltags...')
detector = apriltag.Detector(refine_edges=False)
detections = {}
if visualize:
fig = plt.Figure(figsize=(15, 10))
path = os.path.join(FRAMES, img_paths[0])
img = cv2.imread(path)
image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
result = detector.detect(img)
for i in range(len(result)):
tf = result[i].tag_family
tag_id = result[i].tag_id
cx, cy = result[i].center
# print('Found tag: ', tag_id)
img = cv2.circle(img, (int(cx), int(cy)), 50, (255, 255, 0), thickness=10)
plt.imshow(image)
plt.axis('off')
#plt.savefig('./AprilTag_Detection_%s.jpg' % VIDEOOUT)
plt.show()
tags = dict()
for frame, p in enumerate(img_paths):
path = os.path.join(FRAMES, p)
img = cv2.imread(path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
result = detector.detect(img)
for i in range(len(result)):
tf = result[i].tag_family
cx, cy = result[i].center
tag_id = result[i].tag_id
tags[tag_id] = [cx, cy]
print('Frame %i found %i tags' % (frame, len(result)))
detections[frame] = tags
df = pd.DataFrame.from_dict(detections, orient='index')
path = './AprilTag_%s.pkl' % VIDEOOUT
df.to_pickle(path)
print('Saved AprilTag detections to %s' % path)
if __name__ == '__main__':
process('./Data/ShowCase_3.mp4')