diff options
Diffstat (limited to 'catapult/telemetry/telemetry/internal/image_processing/screen_finder.py')
| -rwxr-xr-x | catapult/telemetry/telemetry/internal/image_processing/screen_finder.py | 857 |
1 files changed, 857 insertions, 0 deletions
diff --git a/catapult/telemetry/telemetry/internal/image_processing/screen_finder.py b/catapult/telemetry/telemetry/internal/image_processing/screen_finder.py new file mode 100755 index 00000000..932d6dfe --- /dev/null +++ b/catapult/telemetry/telemetry/internal/image_processing/screen_finder.py @@ -0,0 +1,857 @@ +#!/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() |