class PreviewWindow(QMainWindow):
"""
QMainWindow subclass used to show frames & tracking.
"""
def __init__(self, controller):
QMainWindow.__init__(self)
# set controller
self.controller = controller
# set title
self.setWindowTitle("Preview")
# get parameter window position & size
param_window_x = self.controller.param_window.x()
param_window_y = self.controller.param_window.y()
param_window_width = self.controller.param_window.width()
# set position & size to be next to the parameter window
self.setGeometry(param_window_x + param_window_width, param_window_y,
10, 10)
# create main widget
self.main_widget = QWidget(self)
self.main_widget.setMinimumSize(QSize(500, 500))
# create main layout
self.main_layout = QGridLayout(self.main_widget)
self.main_layout.setContentsMargins(0, 0, 0, 0)
self.main_layout.setSpacing(0)
# create label that shows frames
self.image_widget = QWidget(self)
self.image_layout = QVBoxLayout(self.image_widget)
self.image_layout.setContentsMargins(0, 0, 0, 0)
self.image_label = PreviewQLabel(self)
self.image_label.setSizePolicy(QSizePolicy.MinimumExpanding,
QSizePolicy.MinimumExpanding)
self.image_label.setAlignment(Qt.AlignTop | Qt.AlignLeft)
self.image_label.hide()
self.image_layout.addWidget(self.image_label)
self.main_layout.addWidget(self.image_widget, 0, 0)
# self.image_label.setStyleSheet("border: 1px solid rgba(122, 127, 130, 0.5)")
self.bottom_widget = QWidget(self)
self.bottom_layout = QVBoxLayout(self.bottom_widget)
self.bottom_layout.setContentsMargins(8, 0, 8, 8)
self.bottom_widget.setMaximumHeight(40)
self.main_layout.addWidget(self.bottom_widget, 1, 0)
# create label that shows crop instructions
self.instructions_label = QLabel("")
self.instructions_label.setStyleSheet("font-size: 11px;")
self.instructions_label.setAlignment(Qt.AlignCenter)
self.bottom_layout.addWidget(self.instructions_label)
# create image slider
self.image_slider = QSlider(Qt.Horizontal)
self.image_slider.setFocusPolicy(Qt.StrongFocus)
self.image_slider.setTickPosition(QSlider.NoTicks)
self.image_slider.setTickInterval(1)
self.image_slider.setSingleStep(1)
self.image_slider.setValue(0)
self.image_slider.valueChanged.connect(self.controller.show_frame)
self.image_slider.hide()
self.bottom_layout.addWidget(self.image_slider)
self.zoom = 1
self.offset = [0, 0]
self.center_y = 0
self.center_x = 0
# initialize variables
self.image = None # image to show
self.tracking_data = None # list of tracking data
self.selecting_crop = False # whether user is selecting a crop
self.changing_heading_angle = False # whether the user is changing the heading angle
self.body_crop = None
self.final_image = None
# set main widget
self.setCentralWidget(self.main_widget)
# set window buttons
if pyqt_version == 5:
self.setWindowFlags(Qt.CustomizeWindowHint
| Qt.WindowMinimizeButtonHint
| Qt.WindowMaximizeButtonHint
| Qt.WindowFullscreenButtonHint)
else:
self.setWindowFlags(Qt.CustomizeWindowHint
| Qt.WindowMinimizeButtonHint
| Qt.WindowMaximizeButtonHint)
self.show()
def wheelEvent(self, event):
old_zoom = self.zoom
self.zoom = max(1, self.zoom + event.pixelDelta().y() / 100)
self.zoom = int(self.zoom * 100) / 100.0
self.update_image_label(self.final_image, zooming=True)
def start_selecting_crop(self):
# start selecting crop
self.selecting_crop = True
# add instruction text
self.instructions_label.setText("Click & drag to select crop area.")
def plot_image(self,
image,
params,
crop_params,
tracking_results,
new_load=False,
new_frame=False,
show_slider=True,
crop_around_body=False):
if image is None:
self.update_image_label(None)
self.image_slider.hide()
self.image_label.hide()
else:
if new_load:
self.image_label.show()
self.remove_tail_start()
if show_slider:
if not self.image_slider.isVisible():
self.image_slider.setValue(0)
self.image_slider.setMaximum(self.controller.n_frames -
1)
self.image_slider.show()
else:
self.image_slider.hide()
max_inititial_size = 500
if image.shape[0] > max_inititial_size:
min_height = max_inititial_size
min_width = max_inititial_size * image.shape[
1] / image.shape[0]
elif image.shape[1] > max_inititial_size:
min_width = max_inititial_size
min_height = max_inititial_size * image.shape[
0] / image.shape[1]
else:
min_height = image.shape[0]
min_width = image.shape[1]
self.main_widget.setMinimumSize(
QSize(min_width, min_height + self.bottom_widget.height()))
# convert to RGB
if len(image.shape) == 2:
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
# update image
self.image = image.copy()
try:
body_crop = params['body_crop']
except:
body_crop = None
try:
tail_start_coords = params['tail_start_coords']
# add tail start point to image
cv2.circle(image, (int(
round(tail_start_coords[1] - crop_params['offset'][1])),
int(
round(tail_start_coords[0] -
crop_params['offset'][0]))), 1,
(180, 180, 50), -1)
except (KeyError, TypeError) as error:
tail_start_coords = None
if tracking_results is not None:
body_position = tracking_results['body_position']
heading_angle = tracking_results['heading_angle']
# add tracking to image
image = tracking.add_tracking_to_frame(image,
tracking_results,
cropped=True)
if body_crop is not None and body_position is not None:
if not crop_around_body:
# copy image
overlay = image.copy()
# draw tail crop overlay
cv2.rectangle(overlay,
(int(body_position[1] - body_crop[1]),
int(body_position[0] - body_crop[0])),
(int(body_position[1] + body_crop[1]),
int(body_position[0] + body_crop[0])),
(242, 242, 65), -1)
# overlay with the original image
cv2.addWeighted(overlay, 0.2, image, 0.8, 0, image)
self.body_crop = None
else:
self.body_crop = body_crop
if crop_around_body:
_, image = tracking.crop_frame_around_body(
image, body_position, params['body_crop'])
self.final_image = image
# update image label
self.update_image_label(
self.final_image,
zoom=(not (crop_around_body and body_position is not None)),
new_load=new_load)
def draw_crop_selection(self, start_crop_coords, end_crop_coords):
if self.selecting_crop and self.image is not None:
# convert image to rgb
if len(self.image.shape) < 3:
image = np.repeat(self.image[:, :, np.newaxis], 3, axis=2)
else:
image = self.image.copy()
# copy image
overlay = image.copy()
# draw crop selection overlay
cv2.rectangle(overlay,
(start_crop_coords[1], start_crop_coords[0]),
(end_crop_coords[1], end_crop_coords[0]),
(255, 51, 0), -1)
# overlay with the original image
cv2.addWeighted(overlay, 0.5, image, 0.5, 0, image)
# update image label
self.update_image_label(image)
def change_offset(self, prev_coords, new_coords):
self.offset[0] -= new_coords[0] - prev_coords[0]
self.offset[1] -= new_coords[1] - prev_coords[1]
self.update_image_label(self.final_image)
def draw_tail_start(self, rel_tail_start_coords):
if self.controller.params['type'] == "headfixed":
# send new tail start coordinates to controller
self.controller.update_tail_start_coords(rel_tail_start_coords)
# clear instructions text
self.instructions_label.setText("")
if self.image is not None:
image = self.image.copy()
cv2.circle(image, (int(round(rel_tail_start_coords[1])),
int(round(rel_tail_start_coords[0]))), 1,
(180, 180, 50), -1)
# update image label
self.update_image_label(image)
def remove_tail_start(self):
self.update_image_label(self.image)
def add_angle_overlay(self, angle):
image = self.image.copy()
image_height = self.image.shape[0]
image_width = self.image.shape[1]
center_y = image_height / 2
center_x = image_width / 2
cv2.arrowedLine(
image,
(int(center_x -
0.3 * image_height * np.sin((angle + 90) * np.pi / 180)),
int(center_y -
0.3 * image_width * np.cos((angle + 90) * np.pi / 180))),
(int(center_x +
0.3 * image_height * np.sin((angle + 90) * np.pi / 180)),
int(center_y +
0.3 * image_width * np.cos((angle + 90) * np.pi / 180))),
(50, 255, 50), 2)
self.update_image_label(image)
def remove_angle_overlay(self):
self.update_image_label(self.image)
def update_image_label(self,
image,
zoom=True,
new_load=False,
zooming=False):
if image is not None and self.zoom != 1 and zoom:
if zooming:
self.offset[0] = min(
max(
0, self.offset[0] + int(
(self.image_label.image.shape[0]) / 2.0) -
int(round((image.shape[0] / self.zoom) / 2.0))),
image.shape[0] - int(round(image.shape[0] / self.zoom)))
self.offset[1] = min(
max(
0, self.offset[1] + int(
(self.image_label.image.shape[1]) / 2.0) -
int(round((image.shape[1] / self.zoom) / 2.0))),
image.shape[1] - int(round(image.shape[1] / self.zoom)))
else:
self.offset[0] = min(
max(0, self.offset[0]),
image.shape[0] - int(round(image.shape[0] / self.zoom)))
self.offset[1] = min(
max(0, self.offset[1]),
image.shape[1] - int(round(image.shape[1] / self.zoom)))
if self.center_y is None:
self.center_y = int(round(image.shape[0] / 2.0))
if self.center_x is None:
self.center_x = int(round(image.shape[1] / 2.0))
image = image[
self.offset[0]:int(round(image.shape[0] / self.zoom)) +
self.offset[0],
self.offset[1]:int(round(image.shape[1] / self.zoom)) +
self.offset[1], :].copy()
if image is not None:
if zoom:
self.setWindowTitle("Preview - Zoom: {:.1f}x".format(
self.zoom))
else:
self.setWindowTitle("Preview - Zoom: 1x")
else:
self.setWindowTitle("Preview")
self.image_label.update_pixmap(image, new_load=new_load)
def crop_selection(self, start_crop_coord, end_crop_coord):
if self.selecting_crop:
# stop selecting the crop
self.selecting_crop = False
# clear instruction text
self.instructions_label.setText("")
# update crop parameters from the selection
self.controller.update_crop_from_selection(start_crop_coord,
end_crop_coord)
def closeEvent(self, ce):
if not self.controller.closing:
ce.ignore()
else:
ce.accept()
