95 lines
2.7 KiB
Python
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')
|