snappysonic/overlay_widget/overlay.py
# coding=utf-8
"""Main loop for tracking visualisation"""
from packaging import version
from PySide2.QtWidgets import QLabel, QWidget
from cv2 import (rectangle, putText, circle, imread)
from numpy import zeros, uint8
from sksurgeryutils.common_overlay_apps import OverlayBaseWidget
from sksurgeryimage.utilities.weisslogo import WeissLogo
from sksurgeryarucotracker import __version__ as arucoversion
from snappysonic.algorithms.algorithms import (configure_tracker,
lookupimage,
check_us_buffer,
get_bg_image_size,
numpy_to_qpixmap)
class OverlayApp(OverlayBaseWidget):
"""
Inherits from OverlayBaseWidget,
adding code to read in video buffers, and display a frame
of data that depends on the position of an external tracking system,
e.g. surgeryarucotracker
"""
def __init__(self, config):
"""
Overides overlay base app's init, to initialise the
external tracking system. Together with a video source
:param config: a configuration dictionary
"""
if "ultrasound buffer" in config:
#and call the constructor for the base class
super().__init__(config.get("ultrasound buffer"),
init_vtk_widget=False)
else:
raise KeyError("Configuration must contain an ultrasound buffer")
self._video_buffers = self._fill_video_buffers(config)
self._tracker = None
if "tracker config" in config:
self._tracker = configure_tracker(config.get("tracker config"))
self._bgimage_offsets = (0, 0)
self._backgroundimage = self._create_background_image(config)
self._weiss = None
self._defaultimage = zeros(0, dtype=uint8)
if "default image" in config:
self._defaultimage = imread(config.get("default image"))
else:
self._weiss = WeissLogo()
self._logger = None
self._tracking_window = QWidget()
self._tracking_window.setWindowTitle("Tracker Position")
self._tracking_viewer = QLabel("Tracking", self._tracking_window)
height, width = self._backgroundimage.shape
self._tracking_window.resize(width, height)
self._tracking_viewer.resize(width, height)
self._tracking_window.show()
self._tracking_viewer.setPixmap(numpy_to_qpixmap(self._backgroundimage))
default_target = 0
if version.parse(arucoversion) >= version.parse("0.2.0"):
default_target = "DICT_4X4_50:0"
self._tracking_target = config.get("tracking target", default_target)
#we could implement something like this?
#if "log directory" in config:
# self._logger = sksurgerydatasaver(config.get("log directory"))
def update_view(self):
"""
Update the background renderer with a new frame,
move the model and render
"""
#add a method to move the rendered models
image = self._get_image_with_tracking()
self.vtk_overlay_window.set_video_image(image)
self.vtk_overlay_window.Render()
def _get_image_with_tracking(self):
"""
Internal method to get an image from the video
buffer based on the tracker position
"""
port_handles, _, _, tracking, _ = self._tracker.get_frame()
tempimg = self._backgroundimage.copy()
pts = None
if port_handles:
for index, port_hdl in enumerate(port_handles):
if port_hdl == self._tracking_target:
pts = (int(tracking[index][0, 3]),
int(tracking[index][1, 3]))
off_pts = (int(pts[0] - self._bgimage_offsets[0]),
int(pts[1] - self._bgimage_offsets[1]))
circle(tempimg, off_pts, 5, [255, 255, 255])
self._tracking_viewer.setPixmap(numpy_to_qpixmap(tempimg))
if pts:
for usbuffer in self._video_buffers:
inframe, image = lookupimage(usbuffer, pts)
if inframe:
return image
if self._defaultimage.shape == (0,):
return self._weiss.get_noisy_logo()
return self._defaultimage
def _fill_video_buffers(self, config):
"""
internal method to fill video buffers
"""
vidbuffers = []
frame_counter = 0
if "buffer descriptions" in config:
for usbuffer in config.get("buffer descriptions"):
check_us_buffer(usbuffer)
start_frame = usbuffer.get("start frame")
end_frame = usbuffer.get("end frame")
tempbuffer = []
if start_frame == frame_counter:
while frame_counter <= end_frame:
ret, image = self.video_source.read()
if not ret:
raise ValueError("Failed Reading video file",
config.get("ultrasound buffer"))
frame_counter = frame_counter + 1
tempbuffer.append(image)
print("adding frame ", len(tempbuffer), " images to buffer ",
usbuffer.get("name"))
usbuffer.update({"buffer" : tempbuffer})
vidbuffers.append(usbuffer)
return vidbuffers
def _create_background_image(self, config):
"""
Creates a backgound image on which we can draw tracking information.
:param config: a configuration dictionary
"""
bg_image_size = [480, 640]
if "buffer descriptions" in config:
self._bgimage_offsets, bg_image_size = get_bg_image_size(config)
bgimage = zeros((bg_image_size), uint8)
if "buffer descriptions" in config:
for usbuffer in config.get("buffer descriptions"):
pt0 = (usbuffer.get("x0") - self._bgimage_offsets[0],
usbuffer.get("y0") - self._bgimage_offsets[1])
pt1 = (usbuffer.get("x1") - self._bgimage_offsets[0],
usbuffer.get("y1") - self._bgimage_offsets[1])
rectangle(bgimage, pt0, pt1, [255, 255, 255])
putText(bgimage, usbuffer.get("name"), pt0, 0,
1.0, [255, 255, 255])
return bgimage