Commit 6c7ba233 authored by Thomas Dowrick's avatar Thomas Dowrick
Browse files

Issue #12: Add AutoCropper class to crop the black background from the davinci xi video feed.

parent 4c82b8a5
Pipeline #1878 failed with stages
in 70 minutes and 51 seconds
......@@ -98,8 +98,13 @@ class UI(QWidget):
def add_crop_buttons(self):
""" Add buttons to crop the video stream. """
self.crop_buttons_layout = QHBoxLayout()
self.crop_button = QPushButton('Crop')
self.other_buttons.addWidget(self.crop_button)
self.autocrop_button = QPushButton('Enable Auto Crop')
self.crop_buttons_layout.addWidget(self.crop_button)
self.crop_buttons_layout.addWidget(self.autocrop_button)
self.other_buttons.addLayout(self.crop_buttons_layout)
def on_exit_clicked(self):
""" Close the application when the exit button has been clicked"""
......
......@@ -2,10 +2,13 @@
import logging
import datetime
from PySide2 import QtWidgets
from PySide2.QtCore import QTimer
from sksurgeryvtk.widgets.vtk_overlay_window import VTKOverlayWindow
from sksurgeryutils.common_overlay_apps import OverlayOnVideoFeedCropRecord
from sksurgeryutils.utils.screen_utils import ScreenController
from sksurgerydavinci.ui.gui import UI
from sksurgerydavinci.widgets.auto_cropping import AutoCropBlackBorder
LOGGER = logging.getLogger(__name__)
......@@ -68,6 +71,11 @@ class MonoViewer(QtWidgets.QWidget):
self.UI.screenshot_button.clicked.connect(self.on_screenshot_clicked)
self.UI.record_button.clicked.connect(self.on_record_start_clicked)
self.UI.crop_button.clicked.connect(self.on_crop_clicked)
self.UI.autocrop_button.clicked.connect(self.on_autocrop_started)
self.auto_cropper = AutoCropBlackBorder(threshold = 75)
self.autocrop_timer = QTimer()
self.autocrop_timer.timeout.connect(self.update_autocrop)
def add_vtk_models(self, models):
"""
......@@ -104,6 +112,28 @@ class MonoViewer(QtWidgets.QWidget):
""" Crop the incoming video stream using ImageCropper. """
self.overlay_window.set_roi()
def on_autocrop_started(self):
""" Start auto cropping. """
self.UI.autocrop_button.setText("Disable Auto Crop")
self.UI.autocrop_button.clicked.disconnect()
self.UI.autocrop_button.clicked.connect(self.on_autocrop_stopped)
self.autocrop_timer.start(500)
def on_autocrop_stopped(self):
""" Stop auto cropping. """
self.UI.autocrop_button.setText("Enable Auto Crop")
self.UI.autocrop_button.clicked.disconnect()
self.UI.autocrop_button.clicked.connect(self.on_autocrop_started)
self.autocrop_timer.stop()
self.overlay_window.roi = None
def update_autocrop(self):
""" Automatically crop the incoming video stream using AutoCropper.
"""
roi = self.auto_cropper.get_roi(self.overlay_window.img)
self.overlay_window.roi = roi
class StereoViewerBase(QtWidgets.QWidget):
"""
......
import numpy as np
import cv2
import time
from sksurgeryimage.acquire.video_source import TimestampedVideoSource
class AutoCropBlackBorder:
""" Class to generate a ROI for an image based on a threshold value.
Originally developed for auto cropping video feed from the DaVinci Xi and
Tile Pro, where there is a central image (which we want to keep), surrounded
by a black region, and then other stuff that we're not interested in.
By searching for the start/end of the black region, we can auto crop the
bit we do want.
For speed, only the red channel of the image is scanned when detecting
black areas (See comment in set_roi for more details).
:param threshold: Threshold value
: param min_size: Minimum size (applies to both width and length)
of the expected output.
"""
def __init__(self, threshold=1, min_size=0):
self.min_size = min_size
self.threshold = threshold
self.roi = None
def get_bounds(self, vector):
""" Given a vector, skip any leading values that are > threshold,
then find the first contiguous range of values that are below the
threshold value.
If the whole vector is above the threshold, return the start/end indexes.
e.g.
[0 0 0 50 50 50 0 0 0 ] return 3, 5 - [50 50 50]
[25 25 0 0 25 25 25 0 0] returns 4, 7 [25 25 25] as the leading 25s are
skipped.
[100 100 100 100 100 100] returns 0, 5.
:param vector: Input vector, in which to find the start/end bounds
:return start: First element that is above the threshold.
:return end: First element following start, that is below threshold.
"""
# If the image starts with a non black area, skip this. We want
# to start searching when the black border begins.
skip = 0
if vector[0] > self.threshold:
skip = np.argmax(vector < self.threshold)
start = skip + np.argmax(vector[skip:] > self.threshold)
end = start + self.min_size + \
np.argmax(vector[(start + self.min_size):] < self.threshold)
# HANDLE EDGE CASES
# argmax returns 0 if there is no match to the search criteria,
# which can give incorrect results, so it is necessary to
# identify some edge cases and act appropriately.
# print(f"skip: {skip}, start: {start}, end: {end}")
# There is a a border at the start, but not the end
if start == end and start != skip:
end = len(vector)
# There is a border at the end, but not at the start
if skip == start == end:
start = 0
# There is no border
if skip == start == end == 0:
end = len(vector)
return start, end
def get_roi(self, img):
""" Calculate the ROI.
Find the x/y extent of the ROI by averaging row and column values,
and then finding the area of interest."""
# Only use the red channel, otherwise we need to take the average along
# two dimensions, which is slow. This should be acceptable for the
# DaVinci.
col_mean = np.mean(img[:,:,1], axis=(0))
row_mean = np.mean(img[:,:,1], axis=(1))
start_y, end_y = self.get_bounds(row_mean)
start_x, end_x = self.get_bounds(col_mean)
self.roi = []
self.roi.append((start_x, start_y))
self.roi.append((end_x, end_y))
return self.roi
if __name__ == "__main__":
input_file = 'tests/data/davinci_xi_generated_video.avi'
vid = cv2.VideoCapture(input_file)
crop = AutoCropBlackBorder()
crop.min_size = 50
crop.threshold = 10
while vid.isOpened():
ret, img = vid.read()
cv2.imshow('orig', img)
roi = crop.get_roi(img)
start_x, start_y = roi[0]
end_x, end_y = roi[1]
cv2.imshow('crop', img[start_y:end_y, start_x:end_x, :])
cv2.waitKey(10)
import cv2
import numpy as np
import pytest
from sksurgerydavinci.widgets.auto_cropping import AutoCropBlackBorder
def set_central_screen(img, size):
"""
Set a square of pixels around the middle of the image
to random pixel values.
"""
height, width = img.shape[:2]
mid_y, mid_x = height // 2, width // 2
dist_from_mid = size // 2
start_x = mid_x - dist_from_mid
end_x = mid_x + dist_from_mid
start_y = mid_y - dist_from_mid
end_y = mid_y + dist_from_mid
# Generate random data to 'put' in the middle of the image
data_size = (end_y - start_y, end_x - start_x, 3)
lower = 0
upper = 255
rand_data = np.random.randint(lower, upper, data_size, dtype = np.uint8)
img[start_y:end_y, start_x:end_x, :] = rand_data
# Return the start/end points of the altered region
roi = [ (start_x, start_y), (end_x, end_y)]
return roi
# def test_roi_all_black_border():
# """ Generated image has black border, then roi e.g.
# """
# width, height = 1920, 1080
# crop = AutoCropBlackBorder()
# for size in [100, 101, 200, 500, 1000, 1001]:
# img = np.zeros((height, width, 3), dtype=np.uint8)
# expected_roi = set_central_screen(img, size)
# calculated_roi = crop.get_roi(img)
# assert expected_roi == calculated_roi
# def test_roi_black_border_first_row_white():
# """ The first row on the DaVinci Xi stream is white (for some reason)
# so make sure we can handle this case.
# """
# width, height = 1920, 1080
# crop = AutoCropBlackBorder()
# for size in [100, 101, 200, 500, 1000, 1001]:
# img = np.zeros((height, width, 3), dtype=np.uint8)
# img[0,:,:] = 255
# expected_roi = set_central_screen(img, size)
# calculated_roi = crop.get_roi(img)
# assert expected_roi == calculated_roi
def test_no_border():
width, height = 1920, 1080
crop = AutoCropBlackBorder()
img = 255 * np.ones((height, width, 3), dtype=np.uint8)
expected_roi = [(0,0), (width, height)]
calculated_roi = crop.get_roi(img)
assert expected_roi == calculated_roi
def test_no_border_bottom_right():
width, height = 1920, 1080
crop = AutoCropBlackBorder()
img = np.zeros((height, width, 3), dtype=np.uint8)
mid_x, mid_y = width // 2, height // 2
img [mid_y:, mid_x:, :] = 255
expected_roi = [(mid_x, mid_y), (width, height)]
calculated_roi = crop.get_roi(img)
assert expected_roi == calculated_roi
def test_no_border_top_left():
width, height = 1920, 1080
crop = AutoCropBlackBorder()
img = np.zeros((height, width, 3), dtype=np.uint8)
mid_x, mid_y = width // 2, height // 2
img [:mid_y, :mid_x, :] = 255
expected_roi = [(0,0), (mid_x, mid_y)]
calculated_roi = crop.get_roi(img)
assert expected_roi == calculated_roi
# start_x, start_y = roi[0]
# end_x, end_y = roi[1]
# print(f"size: {size} w: {end_x-start_x}, h: {end_y-start_y}")
# cv2.imshow('Clipped', img[start_y:end_y,start_x:end_x,:])
# cv2.waitKey(500)
# img[0,:,:] = 255
# roi = crop.get_roi(img)
# start_x, start_y = roi[0]
# end_x, end_y = roi[1]
# print(f"size: {size} w: {end_x-start_x}, h: {end_y-start_y}")
# cv2.imshow('Clipped', img[start_y:end_y,start_x:end_x,:])
# cv2.waitKey(500)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment