| #!/usr/bin/env python |
| # Copyright 2014 The Chromium Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| # |
| # This script attempts to detect the region of a camera's field of view that |
| # contains the screen of the device we are testing. |
| # |
| # Usage: ./screen_finder.py path_to_video 0 0 --verbose |
| |
| from __future__ import division |
| |
| import copy |
| import logging |
| import os |
| import sys |
| |
| if __name__ == '__main__': |
| sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) |
| |
| from telemetry.internal.image_processing import cv_util |
| from telemetry.internal.image_processing import frame_generator as \ |
| frame_generator_module |
| from telemetry.internal.image_processing import video_file_frame_generator |
| from telemetry.internal.util import external_modules |
| |
| np = external_modules.ImportRequiredModule('numpy') |
| cv2 = external_modules.ImportRequiredModule('cv2') |
| |
| |
| class ScreenFinder(object): |
| """Finds and extracts device screens from video. |
| |
| Sample Usage: |
| sf = ScreenFinder(sys.argv[1]) |
| while sf.HasNext(): |
| ret, screen = sf.GetNext() |
| |
| Attributes: |
| _lost_corners: Each index represents whether or not we lost track of that |
| corner on the previous frame. Ordered by [top-right, top-left, |
| bottom-left, bottom-right] |
| _frame: An unmodified copy of the frame we're currently processing. |
| _frame_debug: A copy of the frame we're currently processing, may be |
| modified at any time, used for debugging. |
| _frame_grey: A greyscale copy of the frame we're currently processing. |
| _frame_edges: A Canny Edge detected copy of the frame we're currently |
| processing. |
| _screen_size: The size of device screen in the video when first detected. |
| _avg_corners: Exponentially weighted average of the previous corner |
| locations. |
| _prev_corners: The location of the corners in the previous frame. |
| _lost_corner_frames: A counter of the number of successive frames in which |
| we've lost a corner location. |
| _border: See |border| above. |
| _min_line_length: The minimum length a line must be before we consider it |
| a possible screen edge. |
| _frame_generator: See |frame_generator| above. |
| _width, _height: The width and height of the frame. |
| _anglesp5, _anglesm5: The angles for each point we look at in the grid |
| when computing brightness, constant across frames.""" |
| |
| class ScreenNotFoundError(Exception): |
| pass |
| |
| # Square of the distance a corner can travel in pixels between frames |
| MAX_INTERFRAME_MOTION = 25 |
| # The minimum width line that may be considered a screen edge. |
| MIN_SCREEN_WIDTH = 40 |
| # Number of frames with lost corners before we ignore MAX_INTERFRAME_MOTION |
| RESET_AFTER_N_BAD_FRAMES = 2 |
| # The weight applied to the new screen location when exponentially averaging |
| # screen location. |
| # TODO(mthiesse): This should be framerate dependent, for lower framerates |
| # this value should approach 1. For higher framerates, this value should |
| # approach 0. The current 0.5 value works well in testing with 240 FPS. |
| CORNER_AVERAGE_WEIGHT = 0.5 |
| |
| # TODO(mthiesse): Investigate how to select the constants used here. In very |
| # bright videos, twice as bright may be too high, and the minimum of 60 may |
| # be too low. |
| # The factor by which a quadrant at an intersection must be brighter than |
| # the other quadrants to be considered a screen corner. |
| MIN_RELATIVE_BRIGHTNESS_FACTOR = 1.5 |
| # The minimum average brightness required of an intersection quadrant to |
| # be considered a screen corner (on a scale of 0-255). |
| MIN_CORNER_ABSOLUTE_BRIGHTNESS = 60 |
| |
| # Low and high hysteresis parameters to be passed to the Canny edge |
| # detection algorithm. |
| CANNY_HYSTERESIS_THRESH_LOW = 300 |
| CANNY_HYSTERESIS_THRESH_HIGH = 500 |
| |
| SMALL_ANGLE = 5 / 180 * np.pi # 5 degrees in radians |
| |
| DEBUG = False |
| |
| def __init__(self, frame_generator, border=5): |
| """Initializes the ScreenFinder object. |
| |
| Args: |
| frame_generator: FrameGenerator, An initialized Video Frame Generator. |
| border: int, number of pixels of border to be kept when cropping the |
| detected screen. |
| |
| Raises: |
| FrameReadError: The frame generator may output a read error during |
| initialization.""" |
| assert isinstance(frame_generator, frame_generator_module.FrameGenerator) |
| self._lost_corners = [False, False, False, False] |
| self._frame_debug = None |
| self._frame = None |
| self._frame_grey = None |
| self._frame_edges = None |
| self._screen_size = None |
| self._avg_corners = None |
| self._prev_corners = None |
| self._lost_corner_frames = 0 |
| self._border = border |
| self._min_line_length = self.MIN_SCREEN_WIDTH |
| self._frame_generator = frame_generator |
| self._anglesp5 = None |
| self._anglesm5 = None |
| |
| if not self._InitNextFrame(): |
| logging.warn('Not enough frames in video feed!') |
| return |
| |
| self._height, self._width = self._frame.shape[:2] |
| |
| def _InitNextFrame(self): |
| """Called after processing each frame, reads in the next frame to ensure |
| HasNext() is accurate.""" |
| self._frame_debug = None |
| self._frame = None |
| self._frame_grey = None |
| self._frame_edges = None |
| try: |
| frame = next(self._frame_generator.Generator) |
| except StopIteration: |
| return False |
| self._frame = frame |
| self._frame_debug = copy.copy(frame) |
| self._frame_grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) |
| self._frame_edges = cv2.Canny(self._frame_grey, |
| self.CANNY_HYSTERESIS_THRESH_LOW, |
| self.CANNY_HYSTERESIS_THRESH_HIGH) |
| return True |
| |
| def HasNext(self): |
| """True if there are more frames available to process. """ |
| return self._frame is not None |
| |
| def GetNext(self): |
| """Gets the next screen image. |
| |
| Returns: |
| A numpy matrix containing the screen surrounded by the number of border |
| pixels specified in initialization, and the location of the detected |
| screen corners in the current frame, if a screen is found. The returned |
| screen is guaranteed to be the same size at each frame. |
| 'None' and 'None' if no screen was found on the current frame. |
| |
| Raises: |
| FrameReadError: An error occurred in the FrameGenerator. |
| RuntimeError: This method was called when no frames were available.""" |
| if self._frame is None: |
| raise RuntimeError('No more frames available.') |
| |
| logging.info('Processing frame: %d', |
| self._frame_generator.CurrentFrameNumber) |
| |
| # Finds straight lines in the image. |
| hlines = cv2.HoughLinesP(self._frame_edges, 1, np.pi / 180, 60, |
| minLineLength=self._min_line_length, |
| maxLineGap=100) |
| |
| # Extends these straight lines to be long enough to ensure the screen edge |
| # lines intersect. |
| lines = cv_util.ExtendLines(np.float32(hlines[0]), 10000) \ |
| if hlines is not None else [] |
| |
| # Find intersections in the lines; these are likely to be screen corners. |
| intersections = self._FindIntersections(lines) |
| if len(intersections[:, 0]) > 0: |
| points = np.vstack(intersections[:, 0].flat) |
| if (self._prev_corners is not None and len(points) >= 4 and |
| not self._HasMovedFast(points, self._prev_corners)): |
| corners = self._prev_corners |
| missing_corners = 0 |
| else: |
| # Extract the corners from all intersections. |
| corners, missing_corners = self._FindCorners( |
| intersections, self._frame_grey) |
| else: |
| corners = np.empty((4, 2), np.float32) |
| corners[:] = np.nan |
| missing_corners = 4 |
| |
| screen = None |
| found_screen = True |
| final_corners = None |
| try: |
| # Handle the cases where we have missing corners. |
| screen_corners = self._NewScreenLocation( |
| corners, missing_corners, intersections) |
| |
| final_corners = self._SmoothCorners(screen_corners) |
| |
| # Create a perspective transform from our corners. |
| transform, w, h = self._GetTransform(final_corners, self._border) |
| |
| # Apply the perspective transform to get our output. |
| screen = cv2.warpPerspective( |
| self._frame, transform, (int(w + 0.5), int(h + 0.5))) |
| |
| self._prev_corners = final_corners |
| |
| except self.ScreenNotFoundError as e: |
| found_screen = False |
| logging.info(e) |
| |
| if self.DEBUG: |
| self._Debug(lines, corners, final_corners, screen) |
| |
| self._InitNextFrame() |
| if found_screen: |
| return screen, self._prev_corners |
| return None, None |
| |
| def _FindIntersections(self, lines): |
| """Finds intersections in a set of lines. |
| |
| Filters pairs of lines that are less than 45 degrees apart. Filtering |
| these pairs helps dramatically reduce the number of points we have to |
| process, as these points could not represent screen corners anyways. |
| |
| Returns: |
| The intersections, represented as a tuple of (point, line, line) of the |
| points and the lines that intersect there of all lines in the array that |
| are more than 45 degrees apart.""" |
| intersections = np.empty((0, 3), np.float32) |
| for i in xrange(0, len(lines)): |
| for j in xrange(i + 1, len(lines)): |
| # Filter lines that are less than 45 (or greater than 135) degrees |
| # apart. |
| if not cv_util.AreLinesOrthogonal(lines[i], lines[j], (np.pi / 4.0)): |
| continue |
| ret, point = cv_util.FindLineIntersection(lines[i], lines[j]) |
| point = np.float32(point) |
| if not ret: |
| continue |
| # If we know where the previous corners are, we can also filter |
| # intersections that are too far away from the previous corners to be |
| # where the screen has moved. |
| if self._prev_corners is not None and \ |
| self._lost_corner_frames <= self.RESET_AFTER_N_BAD_FRAMES and \ |
| not self._PointIsCloseToPreviousCorners(point): |
| continue |
| intersections = np.vstack((intersections, |
| np.array((point, lines[i], lines[j])))) |
| return intersections |
| |
| def _PointIsCloseToPreviousCorners(self, point): |
| """True if the point is close to the previous corners.""" |
| max_dist = self.MAX_INTERFRAME_MOTION |
| if cv_util.SqDistance(self._prev_corners[0], point) <= max_dist or \ |
| cv_util.SqDistance(self._prev_corners[1], point) <= max_dist or \ |
| cv_util.SqDistance(self._prev_corners[2], point) <= max_dist or \ |
| cv_util.SqDistance(self._prev_corners[3], point) <= max_dist: |
| return True |
| return False |
| |
| def _HasMovedFast(self, corners, prev_corners): |
| min_dist = np.zeros(4, np.float32) |
| for i in xrange(4): |
| dist = np.min(cv_util.SqDistances(corners, prev_corners[i])) |
| min_dist[i] = dist |
| # 3 corners can move up to one pixel before we consider the screen to have |
| # moved. TODO(mthiesse): Should this be relaxed? Resolution dependent? |
| if np.sum(min_dist) < 3: |
| return False |
| return True |
| |
| class CornerData(object): |
| |
| def __init__(self, corner_index, corner_location, brightness_score, line1, |
| line2): |
| self.corner_index = corner_index |
| self.corner_location = corner_location |
| self.brightness_score = brightness_score |
| self.line1 = line1 |
| self.line2 = line2 |
| |
| def __gt__(self, corner_data2): |
| return self.corner_index > corner_data2.corner_index |
| |
| def __repr__(self): |
| return ('\nCorner index: ' + str(self.corner_index) + |
| ',\nCorner location: ' + str(self.corner_location) + |
| ',\nBrightness score: ' + str(self.brightness_score) + |
| ',\nline1: ' + str(self.line1) + ',\nline2: ' + str(self.line2)) |
| |
| def _FindCorners(self, intersections, grey_frame): |
| """Finds the screen corners in the image. |
| |
| Given the set of intersections in the image, finds the intersections most |
| likely to be corners. |
| |
| Args: |
| intersections: The array of intersections in the image. |
| grey_frame: The greyscale frame we're processing. |
| |
| Returns: |
| An array of length 4 containing the positions of the corners, or nan for |
| each index where a corner could not be found, and a count of the number |
| of missing corners. |
| The corners are ordered as follows: |
| 1 | 0 |
| ----- |
| 2 | 3 |
| Ex. 3 corners are found from a square of width 2 centered at the origin, |
| the output would look like: |
| '[[1, 1], [np.nan, np.nan], [-1, -1], [1, -1]], 1'""" |
| filtered = [] |
| corners = np.empty((0, 2), np.float32) |
| for corner_pos, score, point, line1, line2 in \ |
| self._LooksLikeCorner(intersections, grey_frame): |
| if self.DEBUG: |
| center = (int(point[0] + 0.5), int(point[1] + 0.5)) |
| cv2.circle(self._frame_debug, center, 5, (0, 255, 0), 1) |
| point.resize(1, 2) |
| corners = np.append(corners, point, axis=0) |
| point.resize(2,) |
| corner_data = self.CornerData(corner_pos, point, score, line1, line2) |
| filtered.append(corner_data) |
| |
| # De-duplicate corners because we may have found many false positives, or |
| # near-misses. |
| self._DeDupCorners(filtered, corners) |
| |
| # Strip everything but the corner location. |
| filtered_corners = np.array( |
| [corner_data.corner_location for corner_data in filtered]) |
| corner_indices = [corner_data.corner_index for corner_data in filtered] |
| |
| # If we have found a corner to replace a lost corner, we want to check |
| # that the corner is not erroneous by ensuring it makes a rectangle with |
| # the 3 known good corners. |
| if len(filtered) == 4: |
| for i in xrange(4): |
| point_info = (filtered[i].corner_location, |
| filtered[i].line1, |
| filtered[i].line2) |
| if (self._lost_corners[i] and |
| not self._PointConnectsToCorners(filtered_corners, point_info)): |
| filtered_corners = np.delete(filtered_corners, i, 0) |
| corner_indices = np.delete(corner_indices, i, 0) |
| break |
| |
| # Ensure corners are sorted properly, inserting nans for missing corners. |
| sorted_corners = np.empty((4, 2), np.float32) |
| sorted_corners[:] = np.nan |
| for i in xrange(len(filtered_corners)): |
| sorted_corners[corner_indices[i]] = filtered_corners[i] |
| |
| # From this point on, our corners arrays are guaranteed to have 4 |
| # elements, though some may be nan. |
| |
| # Filter corners that have moved too far from the previous corner if we |
| # are not resetting known corner information. |
| reset_corners = ( |
| (self._lost_corner_frames > self.RESET_AFTER_N_BAD_FRAMES) |
| and len(filtered_corners) == 4) |
| if self._prev_corners is not None and not reset_corners: |
| sqdists = cv_util.SqDistances(self._prev_corners, sorted_corners) |
| for i in xrange(4): |
| if np.isnan(sorted_corners[i][0]): |
| continue |
| if sqdists[i] > self.MAX_INTERFRAME_MOTION: |
| sorted_corners[i] = np.nan |
| |
| real_corners = self._FindExactCorners(sorted_corners) |
| missing_corners = np.count_nonzero(np.isnan(real_corners)) / 2 |
| return real_corners, missing_corners |
| |
| def _LooksLikeCorner(self, intersections, grey_frame): |
| """Finds any intersections of lines that look like a screen corner. |
| |
| Args: |
| intersections: The numpy array of points, and the lines that intersect |
| at the given point. |
| grey_frame: The greyscale frame we're processing. |
| |
| Returns: |
| An array of: The corner location (0-3), the relative brightness score |
| (to be used to de-duplicate corners later), the point, and the lines |
| that make up the intersection, for all intersections that look like a |
| corner.""" |
| points = np.vstack(intersections[:, 0].flat) |
| lines1 = np.vstack(intersections[:, 1].flat) |
| lines2 = np.vstack(intersections[:, 2].flat) |
| # Map the image to four quadrants defined as the regions between each of |
| # the lines that make up the intersection. |
| line1a1 = np.pi - np.arctan2(lines1[:, 1] - points[:, 1], |
| lines1[:, 0] - points[:, 0]) |
| line1a2 = np.pi - np.arctan2(lines1[:, 3] - points[:, 1], |
| lines1[:, 2] - points[:, 0]) |
| line2a1 = np.pi - np.arctan2(lines2[:, 1] - points[:, 1], |
| lines2[:, 0] - points[:, 0]) |
| line2a2 = np.pi - np.arctan2(lines2[:, 3] - points[:, 1], |
| lines2[:, 2] - points[:, 0]) |
| line1a1 = line1a1.reshape(-1, 1) |
| line1a2 = line1a2.reshape(-1, 1) |
| line2a1 = line2a1.reshape(-1, 1) |
| line2a2 = line2a2.reshape(-1, 1) |
| |
| line_angles = np.concatenate((line1a1, line1a2, line2a1, line2a2), axis=1) |
| np.ndarray.sort(line_angles) |
| |
| # TODO(mthiesse): Investigate whether these should scale with image or |
| # screen size. My intuition is that these don't scale with image size, |
| # though they may be affected by image quality and how blurry the corners |
| # are. See stackoverflow.com/q/7765810/ for inspiration. |
| avg_range = 8.0 |
| num_points = 7 |
| |
| points_m_avg = points - avg_range |
| points_p_avg = points + avg_range |
| # Exclude points near frame boundaries. |
| include = np.where((points_m_avg[:, 0] > 0) & (points_m_avg[:, 1] > 0) & |
| (points_p_avg[:, 0] < self._width) & |
| (points_p_avg[:, 1] < self._height)) |
| line_angles = line_angles[include] |
| points = points[include] |
| lines1 = lines1[include] |
| lines2 = lines2[include] |
| points_m_avg = points_m_avg[include] |
| points_p_avg = points_p_avg[include] |
| # Perform a 2-d linspace to generate the x, y ranges for each |
| # intersection. |
| arr1 = points_m_avg[:, 0].reshape(-1, 1) |
| arr2 = points_p_avg[:, 0].reshape(-1, 1) |
| lin = np.linspace(0, 1, num_points) |
| x_range = arr1 + (arr2 - arr1) * lin |
| arr1 = points_m_avg[:, 1].reshape(-1, 1) |
| arr2 = points_p_avg[:, 1].reshape(-1, 1) |
| y_range = arr1 + (arr2 - arr1) * lin |
| |
| # The angles for each point we look at in the grid when computing |
| # brightness are constant across frames, so we can generate them once. |
| if self._anglesp5 is None: |
| ind = np.transpose([np.tile(x_range[0], num_points), |
| np.repeat(y_range[0], num_points)]) |
| vectors = ind - points[0] |
| angles = np.arctan2(vectors[:, 1], vectors[:, 0]) + np.pi |
| self._anglesp5 = angles + self.SMALL_ANGLE |
| self._anglesm5 = angles - self.SMALL_ANGLE |
| results = [] |
| for i in xrange(len(y_range)): |
| # Generate our filters for which points belong to which quadrant. |
| one = np.where((self._anglesp5 <= line_angles[i, 1]) & |
| (self._anglesm5 >= line_angles[i, 0])) |
| two = np.where((self._anglesp5 <= line_angles[i, 2]) & |
| (self._anglesm5 >= line_angles[i, 1])) |
| thr = np.where((self._anglesp5 <= line_angles[i, 3]) & |
| (self._anglesm5 >= line_angles[i, 2])) |
| fou = np.where((self._anglesp5 <= line_angles[i, 0]) | |
| (self._anglesm5 >= line_angles[i, 3])) |
| # Take the cartesian product of our x and y ranges to get the full list |
| # of pixels to look at. |
| ind = np.transpose([np.tile(x_range[i], num_points), |
| np.repeat(y_range[i], num_points)]) |
| |
| # Filter the full list by which indices belong to which quadrant, and |
| # convert to integers so we can index with them. |
| one_i = np.int32(np.rint(ind[one[0]])) |
| two_i = np.int32(np.rint(ind[two[0]])) |
| thr_i = np.int32(np.rint(ind[thr[0]])) |
| fou_i = np.int32(np.rint(ind[fou[0]])) |
| |
| # Average the brightness of the pixels that belong to each quadrant. |
| q_1 = np.average(grey_frame[one_i[:, 1], one_i[:, 0]]) |
| q_2 = np.average(grey_frame[two_i[:, 1], two_i[:, 0]]) |
| q_3 = np.average(grey_frame[thr_i[:, 1], thr_i[:, 0]]) |
| q_4 = np.average(grey_frame[fou_i[:, 1], fou_i[:, 0]]) |
| |
| avg_intensity = [(q_4, 0), (q_1, 1), (q_2, 2), (q_3, 3)] |
| # Sort by intensity. |
| avg_intensity.sort(reverse=True) |
| |
| # Treat the point as a corner if one quadrant is at least twice as |
| # bright as the next brightest quadrant, with a minimum brightness |
| # requirement. |
| tau = (2.0 * np.pi) |
| min_factor = self.MIN_RELATIVE_BRIGHTNESS_FACTOR |
| min_brightness = self.MIN_RELATIVE_BRIGHTNESS_FACTOR |
| if avg_intensity[0][0] > avg_intensity[1][0] * min_factor and \ |
| avg_intensity[0][0] > min_brightness: |
| bright_corner = avg_intensity[0][1] |
| if bright_corner == 0: |
| angle = np.pi - (line_angles[i, 0] + line_angles[i, 3]) / 2.0 |
| if angle < 0: |
| angle = angle + tau |
| else: |
| angle = tau - (line_angles[i, bright_corner] + |
| line_angles[i, bright_corner - 1]) / 2.0 |
| score = avg_intensity[0][0] - avg_intensity[1][0] |
| # TODO(mthiesse): int(angle / (pi / 2.0)) will break if the screen is |
| # rotated through 45 degrees. Probably many other things will break as |
| # well, movement of corners from one quadrant to another hasn't been |
| # tested. We should support this eventually, but this is unlikely to |
| # cause issues for any test setups. |
| results.append((int(angle / (np.pi / 2.0)), score, points[i], |
| lines1[i], lines2[i])) |
| return results |
| |
| def _DeDupCorners(self, corner_data, corners): |
| """De-duplicate corners based on corner_index. |
| |
| For each set of points representing a corner: If one point is part of the |
| rectangle and the other is not, filter the other one. If both or none are |
| part of the rectangle, filter based on score (highest relative brightness |
| of a quadrant). The reason we allow for neither to be part of the |
| rectangle is because we may not have found all four corners of the |
| rectangle, and in degenerate cases like this it's better to find 3 likely |
| corners than none. |
| |
| Modifies corner_data directly. |
| |
| Args: |
| corner_data: CornerData for each potential corner in the frame. |
| corners: List of all potential corners in the frame.""" |
| # TODO(mthiesse): Ensure that the corners form a sensible rectangle. For |
| # example, it is currently possible (but unlikely) to detect a 'screen' |
| # where the bottom-left corner is above the top-left corner, while the |
| # bottom-right corner is below the top-right corner. |
| |
| # Sort by corner_index to make de-duping easier. |
| corner_data.sort() |
| |
| # De-dup corners. |
| c_old = None |
| for i in xrange(len(corner_data) - 1, 0, -1): |
| if corner_data[i].corner_index != corner_data[i - 1].corner_index: |
| c_old = None |
| continue |
| if c_old is None: |
| point_info = (corner_data[i].corner_location, |
| corner_data[i].line1, |
| corner_data[i].line2) |
| c_old = self._PointConnectsToCorners(corners, point_info, 2) |
| point_info_new = (corner_data[i - 1].corner_location, |
| corner_data[i - 1].line1, |
| corner_data[i - 1].line2) |
| c_new = self._PointConnectsToCorners(corners, point_info_new, 2) |
| if (not (c_old or c_new)) or (c_old and c_new): |
| if (corner_data[i].brightness_score < |
| corner_data[i - 1].brightness_score): |
| del corner_data[i] |
| c_old = c_new |
| else: |
| del corner_data[i - 1] |
| elif c_old: |
| del corner_data[i - 1] |
| else: |
| del corner_data[i] |
| c_old = c_new |
| |
| def _PointConnectsToCorners(self, corners, point_info, tolerance=1): |
| """Checks if the lines of an intersection intersect with corners. |
| |
| This is useful to check if the point is part of a rectangle specified by |
| |corners|. |
| |
| Args: |
| point_info: A tuple of (point, line, line) representing an intersection |
| of two lines. |
| corners: corners that (hopefully) make up a rectangle. |
| tolerance: The tolerance (approximately in pixels) of the distance |
| between the corners and the lines for detecting if the point is on |
| the line. |
| |
| Returns: |
| True if each of the two lines that make up the intersection where the |
| point is located connect the point to other corners.""" |
| line1_connected = False |
| line2_connected = False |
| point, line1, line2 = point_info |
| for corner in corners: |
| if corner is None: |
| continue |
| |
| # Filter out points that are too close to one another to be different |
| # corners. |
| sqdist = cv_util.SqDistance(corner, point) |
| if sqdist < self.MIN_SCREEN_WIDTH * self.MIN_SCREEN_WIDTH: |
| continue |
| |
| line1_connected = line1_connected or \ |
| cv_util.IsPointApproxOnLine(corner, line1, tolerance) |
| line2_connected = line2_connected or \ |
| cv_util.IsPointApproxOnLine(corner, line2, tolerance) |
| if line1_connected and line2_connected: |
| return True |
| return False |
| |
| def _FindExactCorners(self, sorted_corners): |
| """Attempts to find more accurate corner locations. |
| |
| Args: |
| sorted_corners: The four screen corners, sorted by corner_index. |
| |
| Returns: |
| A list of 4 probably more accurate corners, still sorted.""" |
| real_corners = np.empty((4, 2), np.float32) |
| # Count missing corners, and search in a small area around our |
| # intersections representing corners to see if we can find a more exact |
| # corner, as the position of the intersections is noisy and not always |
| # perfectly accurate. |
| for i in xrange(4): |
| corner = sorted_corners[i] |
| if np.isnan(corner[0]): |
| real_corners[i] = np.nan |
| continue |
| |
| # Almost unbelievably, in edge cases with floating point error, the |
| # width/height of the cropped corner image may be 2 or 4. This is fine |
| # though, as long as the width and height of the cropped corner are not |
| # hard-coded anywhere. |
| corner_image = self._frame_edges[corner[1] - 1:corner[1] + 2, |
| corner[0] - 1:corner[0] + 2] |
| ret, p = self._FindExactCorner(i <= 1, i == 1 or i == 2, corner_image) |
| if ret: |
| if self.DEBUG: |
| self._frame_edges[corner[1] - 1 + p[1]][corner[0] - 1 + p[0]] = 128 |
| real_corners[i] = corner - 1 + p |
| else: |
| real_corners[i] = corner |
| return real_corners |
| |
| def _FindExactCorner(self, top, left, img): |
| """Tries to finds the exact corner location for a given corner. |
| |
| Searches for the top or bottom, left or right most lit |
| pixel in an edge-detected image, which should represent, with pixel |
| precision, as accurate a corner location as possible. (Though perhaps |
| up-sampling using cubic spline interpolation could get sub-pixel |
| precision) |
| |
| TODO(mthiesse): This algorithm could be improved by including a larger |
| region to search in, but would have to be made smarter about which lit |
| pixels are on the detected screen edge and which are a not as it's |
| currently extremely easy to fool by things like notification icons in |
| screen corners. |
| |
| Args: |
| top: boolean, whether or not we're looking for a top corner. |
| left: boolean, whether or not we're looking for a left corner. |
| img: A small cropping of the edge detected image in which to search. |
| |
| Returns: |
| True and the location if a better corner location is found, |
| False otherwise.""" |
| h, w = img.shape[:2] |
| cy = 0 |
| starting_x = w - 1 if left else 0 |
| cx = starting_x |
| if top: |
| y_range = xrange(h - 1, -1, -1) |
| else: |
| y_range = xrange(0, h, 1) |
| if left: |
| x_range = xrange(w - 1, -1, -1) |
| else: |
| x_range = xrange(0, w, 1) |
| for y in y_range: |
| for x in x_range: |
| if img[y][x] == 255: |
| cy = y |
| if (left and x <= cx) or (not left and x >= cx): |
| cx = x |
| if cx == starting_x and cy == 0 and img[0][starting_x] != 255: |
| return False, (0, 0) |
| return True, (cx, cy) |
| |
| def _NewScreenLocation(self, new_corners, missing_corners, intersections): |
| """Computes the new screen location with best effort. |
| |
| Creates the final list of corners that represents the best effort attempt |
| to find the new screen location. Handles degenerate cases where 3 or fewer |
| new corners are present, using previous corner and intersection data. |
| |
| Args: |
| new_corners: The corners found by our search for corners. |
| missing_corners: The count of how many corners we're missing. |
| intersections: The intersections of straight lines found in the current |
| frame. |
| |
| Returns: |
| An array of 4 new_corners hopefully representing the screen, or throws |
| an error if this is not possible. |
| |
| Raises: |
| ValueError: Finding the screen location was not possible.""" |
| screen_corners = copy.copy(new_corners) |
| if missing_corners == 0: |
| self._lost_corner_frames = 0 |
| self._lost_corners = [False, False, False, False] |
| return screen_corners |
| if self._prev_corners is None: |
| raise self.ScreenNotFoundError( |
| 'Could not locate screen on frame %d' % |
| self._frame_generator.CurrentFrameNumber) |
| |
| self._lost_corner_frames += 1 |
| if missing_corners > 1: |
| logging.info('Unable to properly detect screen corners, making ' |
| 'potentially false assumptions on frame %d', |
| self._frame_generator.CurrentFrameNumber) |
| # Replace missing new_corners with either nearest intersection to previous |
| # corner, or previous corner if no intersections are found. |
| for i in xrange(0, 4): |
| if not np.isnan(new_corners[i][0]): |
| self._lost_corners[i] = False |
| continue |
| self._lost_corners[i] = True |
| min_dist = self.MAX_INTERFRAME_MOTION |
| min_corner = None |
| |
| for isection in intersections: |
| dist = cv_util.SqDistance(isection[0], self._prev_corners[i]) |
| if dist >= min_dist: |
| continue |
| if missing_corners == 1: |
| # We know in this case that we have 3 corners present, meaning |
| # all 4 screen lines, and therefore intersections near screen |
| # corners present, so our new corner must connect to these |
| # other corners. |
| if not self._PointConnectsToCorners(new_corners, isection, 3): |
| continue |
| min_corner = isection[0] |
| min_dist = dist |
| screen_corners[i] = min_corner if min_corner is not None else \ |
| self._prev_corners[i] |
| |
| return screen_corners |
| |
| def _SmoothCorners(self, corners): |
| """Smoothes the motion of corners, reduces noise. |
| |
| Smoothes the motion of corners by computing an exponentially weighted |
| moving average of corner positions over time. |
| |
| Args: |
| corners: The corners of the detected screen. |
| |
| Returns: |
| The final corner positions.""" |
| if self._avg_corners is None: |
| self._avg_corners = np.asfarray(corners, np.float32) |
| for i in xrange(0, 4): |
| # Keep an exponential moving average of the corner location to reduce |
| # noise. |
| new_contrib = np.multiply(self.CORNER_AVERAGE_WEIGHT, corners[i]) |
| old_contrib = np.multiply(1 - self.CORNER_AVERAGE_WEIGHT, |
| self._avg_corners[i]) |
| self._avg_corners[i] = np.add(new_contrib, old_contrib) |
| |
| return self._avg_corners |
| |
| def _GetTransform(self, corners, border): |
| """Gets the perspective transform of the screen. |
| |
| Args: |
| corners: The corners of the detected screen. |
| border: The number of pixels of border to crop along with the screen. |
| |
| Returns: |
| A perspective transform and the width and height of the target |
| transform. |
| |
| Raises: |
| ScreenNotFoundError: Something went wrong in detecting the screen.""" |
| if self._screen_size is None: |
| w = np.sqrt(cv_util.SqDistance(corners[1], corners[0])) |
| h = np.sqrt(cv_util.SqDistance(corners[1], corners[2])) |
| if w < 1 or h < 1: |
| raise self.ScreenNotFoundError( |
| 'Screen detected improperly (bad corners)') |
| if min(w, h) < self.MIN_SCREEN_WIDTH: |
| raise self.ScreenNotFoundError('Detected screen was too small.') |
| |
| self._screen_size = (w, h) |
| # Extend min line length, if we can, to reduce the number of extraneous |
| # lines the line finder finds. |
| self._min_line_length = max(self._min_line_length, min(w, h) / 1.75) |
| w = self._screen_size[0] |
| h = self._screen_size[1] |
| |
| target = np.zeros((4, 2), np.float32) |
| width = w + border |
| height = h + border |
| target[0] = np.asfarray((width, border)) |
| target[1] = np.asfarray((border, border)) |
| target[2] = np.asfarray((border, height)) |
| target[3] = np.asfarray((width, height)) |
| transform_w = width + border |
| transform_h = height + border |
| transform = cv2.getPerspectiveTransform(corners, target) |
| return transform, transform_w, transform_h |
| |
| def _Debug(self, lines, corners, final_corners, screen): |
| for line in lines: |
| intline = ((int(line[0]), int(line[1])), |
| (int(line[2]), int(line[3]))) |
| cv2.line(self._frame_debug, intline[0], intline[1], (0, 0, 255), 1) |
| i = 0 |
| for corner in corners: |
| if not np.isnan(corner[0]): |
| cv2.putText( |
| self._frame_debug, str(i), (int(corner[0]), int(corner[1])), |
| cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (255, 255, 0), 1, cv2.CV_AA) |
| i += 1 |
| if final_corners is not None: |
| for corner in final_corners: |
| cv2.circle(self._frame_debug, |
| (int(corner[0]), int(corner[1])), 5, (255, 0, 255), 1) |
| cv2.imshow('original', self._frame) |
| cv2.imshow('debug', self._frame_debug) |
| if screen is not None: |
| cv2.imshow('screen', screen) |
| cv2.waitKey() |
| |
| # For being run as a script. |
| # TODO(mthiesse): To be replaced with a better standalone script. |
| # Ex: ./screen_finder.py path_to_video 0 5 --verbose |
| |
| |
| def main(): |
| start_frame = int(sys.argv[2]) if len(sys.argv) >= 3 else 0 |
| vf = video_file_frame_generator.VideoFileFrameGenerator(sys.argv[1], |
| start_frame) |
| if len(sys.argv) >= 4: |
| sf = ScreenFinder(vf, int(sys.argv[3])) |
| else: |
| sf = ScreenFinder(vf) |
| # TODO(mthiesse): Use argument parser to improve command line parsing. |
| if len(sys.argv) > 4 and sys.argv[4] == '--verbose': |
| logging.basicConfig(format='%(message)s', level=logging.INFO) |
| else: |
| logging.basicConfig(format='%(message)s', level=logging.WARN) |
| while sf.HasNext(): |
| sf.GetNext() |
| |
| if __name__ == '__main__': |
| main() |