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

sksurgerydavinci/ui/gui.py

@@ -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"""

sksurgerydavinci/widgets/Viewers.py

@@ -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):
 
     """

sksurgerydavinci/widgets/auto_cropping.py

+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)
+

tests/test_auto_crop.py

+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)
+