# -*- coding: utf-8 -*-
""" Containers for video calibration data. """
import os
import copy
import cv2
import numpy as np
import sksurgerycalibration.video.video_calibration_io as sksio
[docs]class BaseVideoCalibrationData:
"""
Constructor, no member variables, so just a pure virtual interface.
Not really necessary if you rely on duck-typing, but at least
it shows the intention of what derived classes should implement,
and means we can use this base class to type check against.
"""
def __init__(self):
return
[docs] def reinit(self):
""" Used to clear, re-initialise all member variables. """
raise NotImplementedError("Derived classes should implement this.")
[docs] def pop(self):
""" Remove the last view of data. """
raise NotImplementedError("Derived classes should implement this.")
[docs] def get_number_of_views(self):
""" Returns the number of views of data. """
raise NotImplementedError("Derived classes should implement this.")
[docs] def save_data(self, dir_name: str, file_prefix: str):
""" Writes all contained data to disk. """
raise NotImplementedError("Derived classes should implement this.")
[docs] def load_data(self, dir_name: str, file_prefix: str):
""" Loads all contained data from disk. """
raise NotImplementedError("Derived classes should implement this.")
[docs]class TrackingData(BaseVideoCalibrationData):
"""
Class for storing tracking data.
"""
def __init__(self):
super().__init__()
self.device_tracking_array = None
self.calibration_tracking_array = None
self.reinit()
[docs] def reinit(self):
"""
Deletes all data.
"""
self.device_tracking_array = []
self.calibration_tracking_array = []
[docs] def push(self, device_tracking, calibration_tracking):
"""
Stores a pair of tracking data.
:param device_tracking: transformation for the thing you're tracking
:param calibration_tracking: transformation for tracked calibration obj
"""
self.device_tracking_array.append(
copy.deepcopy(device_tracking))
self.calibration_tracking_array.append(
copy.deepcopy(calibration_tracking))
[docs] def pop(self):
"""
Removes the last (most recent) view of data.
"""
if self.device_tracking_array:
self.device_tracking_array.pop(-1)
self.calibration_tracking_array.pop(-1)
[docs] def get_number_of_views(self):
"""
Returns the number of views of data.
:return: int
"""
return len(self.device_tracking_array)
[docs] def save_data(self,
dir_name: str,
file_prefix: str
):
"""
Saves the tracking data to lots of different files.
:param dir_name: directory to save to
:param file_prefix: prefix for all files
"""
if not os.path.isdir(dir_name):
os.makedirs(dir_name)
for i in enumerate(self.device_tracking_array):
device_tracking_file = \
sksio.get_device_tracking_file_name(dir_name,
file_prefix,
i[0])
if self.device_tracking_array[i[0]] is not None:
np.savetxt(device_tracking_file,
self.device_tracking_array[i[0]], fmt='%.8f')
for i in enumerate(self.calibration_tracking_array):
calibration_tracking_file = \
sksio.get_calibration_tracking_file_name(dir_name,
file_prefix,
i[0])
if self.calibration_tracking_array[i[0]] is not None:
np.savetxt(calibration_tracking_file,
self.calibration_tracking_array[i[0]], fmt='%.8f')
[docs] def load_data(self,
dir_name: str,
file_prefix: str
):
"""
Loads tracking data from files.
:param dir_name: directory to load from
:param file_prefix: prefix for all files
"""
self.reinit()
files = sksio.get_filenames_by_glob_expr(dir_name,
file_prefix,
"device_tracking",
".txt")
for file in files:
device_data = np.loadtxt(file)
self.device_tracking_array.append(device_data)
files = sksio.get_filenames_by_glob_expr(dir_name,
file_prefix,
"calib_obj_tracking",
".txt")
for file in files:
calibration_data = np.loadtxt(file)
self.calibration_tracking_array.append(calibration_data)
[docs]class MonoVideoData(BaseVideoCalibrationData):
"""
Stores data extracted from each video view of a mono calibration.
"""
def __init__(self):
super().__init__()
self.images_array = None
self.ids_arrays = None
self.object_points_arrays = None
self.image_points_arrays = None
self.reinit()
[docs] def reinit(self):
"""
Deletes all data.
"""
self.images_array = []
self.ids_arrays = []
self.object_points_arrays = []
self.image_points_arrays = []
[docs] def pop(self):
"""
Removes the last (most recent) view of data.
"""
if len(self.images_array) > 1:
self.images_array.pop(-1)
self.ids_arrays.pop(-1)
self.object_points_arrays.pop(-1)
self.image_points_arrays.pop(-1)
[docs] def push(self, image, ids, object_points, image_points):
"""
Stores another view of data. Copies data.
"""
self.images_array.append(copy.deepcopy(image))
self.ids_arrays.append(copy.deepcopy(ids))
self.object_points_arrays.append(copy.deepcopy(object_points))
self.image_points_arrays.append(copy.deepcopy(image_points))
[docs] def get_number_of_views(self):
"""
Returns the number of views.
"""
return len(self.images_array)
[docs] def save_data(self,
dir_name: str,
file_prefix: str
):
"""
Saves the calibration data to lots of different files.
:param dir_name: directory to save to
:param file_prefix: prefix for all files
"""
if not os.path.isdir(dir_name):
os.makedirs(dir_name)
for i in enumerate(self.images_array):
image_file = sksio.get_images_file_name(dir_name,
file_prefix,
i[0])
cv2.imwrite(image_file, self.images_array[i[0]])
for i in enumerate(self.ids_arrays):
id_file = sksio.get_ids_file_name(dir_name,
file_prefix,
i[0])
np.savetxt(id_file, self.ids_arrays[i[0]], fmt='%d')
for i in enumerate(self.object_points_arrays):
object_points_file = sksio.get_objectpoints_file_name(dir_name,
file_prefix,
i[0])
reshaped = np.reshape(self.object_points_arrays[i[0]], (-1, 3))
np.savetxt(object_points_file, reshaped, fmt='%.8f')
for i in enumerate(self.image_points_arrays):
image_points_file = sksio.get_imagepoints_file_name(dir_name,
file_prefix,
i[0])
reshaped = np.reshape(self.image_points_arrays[i[0]], (-1, 2))
np.savetxt(image_points_file, reshaped, fmt='%.8f')
[docs] def load_data(self,
dir_name: str,
file_prefix: str
):
"""
Loads the calibration data.
:param dir_name: directory to load from
:param file_prefix: prefix for all files
"""
self.reinit()
files = sksio.get_filenames_by_glob_expr(dir_name,
file_prefix,
"images",
".png")
for file in files:
image = cv2.imread(file)
self.images_array.append(image)
files = sksio.get_filenames_by_glob_expr(dir_name,
file_prefix,
"ids",
".txt")
for file in files:
ids = np.loadtxt(file)
self.ids_arrays.append(ids)
files = sksio.get_filenames_by_glob_expr(dir_name,
file_prefix,
"object_points",
".txt")
for file in files:
object_points = np.loadtxt(file)
reshaped = np.reshape(object_points, (object_points.shape[0], 1, 3))
self.object_points_arrays.append(reshaped.astype(np.float32))
files = sksio.get_filenames_by_glob_expr(dir_name,
file_prefix,
"image_points",
".txt")
for file in files:
image_points = np.loadtxt(file)
reshaped = np.reshape(image_points, (image_points.shape[0], 1, 2))
self.image_points_arrays.append(reshaped.astype(np.float32))
[docs] def save_annotated_images(self,
dir_name: str,
file_prefix: str
):
"""
Saves video images, annotated with the ID of
each 2D point detected.
"""
if len(self.images_array) == 0:
raise ValueError("Can't dump out empty arrays")
if not os.path.isdir(dir_name):
os.makedirs(dir_name)
font = cv2.FONT_HERSHEY_SIMPLEX
for i in enumerate(self.images_array):
image_copy = copy.deepcopy(self.images_array[i])
points = self.image_points_arrays[i]
ids = self.ids_arrays[i]
for counter in range(ids.shape[0]):
cv2.putText(image_copy,
str(ids[counter]),
(int(points[counter][0][0]),
int(points[counter][0][1])),
font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
annotated_image_file_name = \
sksio.get_annotated_images_file_name(dir_name,
file_prefix,
i)
cv2.imwrite(annotated_image_file_name, image_copy)
[docs]class StereoVideoData(BaseVideoCalibrationData):
"""
Stores data extracted from each view of a stereo calibration.
"""
def __init__(self):
super().__init__()
self.left_data = MonoVideoData()
self.right_data = MonoVideoData()
self.reinit()
[docs] def reinit(self):
"""
Deletes all data.
"""
self.left_data.reinit()
self.right_data.reinit()
[docs] def pop(self):
"""
Removes the last (most recent) view of data.
"""
self.left_data.pop()
self.right_data.pop()
[docs] def push(self,
left_image, left_ids, left_object_points, left_image_points,
right_image, right_ids, right_object_points, right_image_points):
"""
Stores another view of data. Copies data.
"""
self.left_data.push(
left_image, left_ids, left_object_points, left_image_points)
self.right_data.push(
right_image, right_ids, right_object_points, right_image_points)
[docs] def get_number_of_views(self):
"""
Returns the number of views.
"""
num_left = self.left_data.get_number_of_views()
num_right = self.right_data.get_number_of_views()
if num_left != num_right:
raise ValueError("Different number of views in left and right??")
return num_left
[docs] def save_data(self,
dir_name: str,
file_prefix: str
):
"""
Saves the calibration data to lots of different files.
:param dir_name: directory to save to
:param file_prefix: prefix for all files
"""
left_prefix = sksio.get_left_prefix(file_prefix)
self.left_data.save_data(dir_name, left_prefix)
right_prefix = sksio.get_right_prefix(file_prefix)
self.right_data.save_data(dir_name, right_prefix)
[docs] def load_data(self,
dir_name: str,
file_prefix: str
):
"""
Loads the calibration data.
:param dir_name: directory to load from
:param file_prefix: prefix for all files
"""
self.reinit()
left_prefix = sksio.get_left_prefix(file_prefix)
self.left_data.load_data(dir_name, left_prefix)
right_prefix = sksio.get_right_prefix(file_prefix)
self.right_data.load_data(dir_name, right_prefix)
[docs] def save_annotated_images(self,
dir_name: str,
file_prefix: str
):
"""
Saves video images, annotated with the ID of
each 2D point detected.
"""
left_prefix = sksio.get_left_prefix(file_prefix)
self.left_data.save_annotated_images(dir_name, left_prefix)
right_prefix = sksio.get_right_prefix(file_prefix)
self.right_data.save_annotated_images(dir_name, right_prefix)