diff --git a/variations/variations.py b/variations/variations.py index e2848f0..ff8f967 100644 --- a/variations/variations.py +++ b/variations/variations.py @@ -14,6 +14,8 @@ from tesserocr import PSM, PyTessBaseAPI flags.DEFINE_integer('device_id', None, 'ID of the capture card device.') +flags.DEFINE_boolean('for_sale', False, 'Whether to only process items that are for sale.') +flags.DEFINE_string('video_path', None, 'Path to video file to use instead of capture device.') FLAGS = flags.FLAGS TUTORIAL_LINES = [ @@ -63,7 +65,7 @@ def annotate_frame(self, frame: numpy.ndarray) -> numpy.ndarray: """Parses various parts of a frame for catalog items and annotates it.""" # Detect whether we are in in Nook Shopping catalog. - if not numpy.array_equal(frame[500, 20], (182, 249, 255)): + if numpy.linalg.norm(frame[500, 20] - (182, 249, 255)) > 5: text = 'Navigate to Nook catalog to start!' opts = {'org': (200, 70), 'fontFace': cv2.FONT_HERSHEY_PLAIN, 'lineType': cv2.LINE_AA, 'fontScale': 3} @@ -84,8 +86,8 @@ def annotate_frame(self, frame: numpy.ndarray) -> numpy.ndarray: frame = cv2.putText(frame, count_text, (500, 700), 0, 1, 0) gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - section = numpy.nonzero(gray[20, 250:] == 156)[0] - if section.any() and section[0] != self.active_section: + section = numpy.where((154 < gray[20, 250:]) & (gray[20, 250:] < 162))[0] + if section.any() and abs(self.active_section - section[0]) > 10: # Grab the new section name x1, *_, x2 = 250 + section section_region = 255 - gray[8:32, x1+5:x2-5] @@ -131,7 +133,7 @@ def annotate_frame(self, frame: numpy.ndarray) -> numpy.ndarray: def get_selected_item(self, frame: numpy.ndarray) -> Optional[Rectangle]: """Returns the rectangle around the selected item name if there is one.""" # Search for the yellow selected region along the item list area. - select_region = numpy.nonzero(frame[140:640, 1052, 0] < 100)[0] + select_region = numpy.where(frame[140:640, 1052, 0] < 100)[0] if not select_region.any(): return None @@ -147,7 +149,7 @@ def get_selected_item(self, frame: numpy.ndarray) -> Optional[Rectangle]: # Detect to width of the name by collapsing along the x axis # and finding the right-most dark pixel (text). item_region = frame[rect.y1:rect.y2, rect.x1:rect.x2, 1] - detected_text = numpy.nonzero(item_region.min(axis=0) < 55)[0] + detected_text = numpy.where(item_region.min(axis=0) < 55)[0] if not detected_text.any(): return None @@ -157,7 +159,7 @@ def get_selected_item(self, frame: numpy.ndarray) -> Optional[Rectangle]: def get_variation(self, gray: numpy.ndarray) -> Optional[Rectangle]: """Returns the rectangle around the variation text if there is one.""" # There's a white box if the item has a variation. - if gray[650, 25] != 250: + if gray[650, 25] < 250: return None variation = Rectangle(x1=30, y1=628, y2=665) @@ -196,6 +198,7 @@ def save_items(self) -> None: date = datetime.datetime.now().strftime('%d-%m-%Y %H-%M-%S') with open(f'{self.section_name} ({date}).txt', 'w') as fp: + print(f'Saving {len(self.items)} items to {fp.name}') fp.write('\n'.join(sorted(self.items))) self.section_name = None @@ -243,22 +246,29 @@ def main(argv): os.chdir(data_dir) # Find the right device ID. - device_id = FLAGS.device_id - if device_id is None: - device_id = pick_device_id() - - # Connect to the capture card and adjust video dimensions. - cap = cv2.VideoCapture(device_id) + if FLAGS.video_path: + video_capture = FLAGS.video_path + print(f'Processing video file {video_capture}') + elif FLAGS.device_id is not None: + video_capture = FLAGS.device_id + else: + video_capture = pick_device_id() + + # Connect to input feed and adjust video dimensions. + cap = cv2.VideoCapture(video_capture) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) - cv2.namedWindow('frame') parser = VariationParser() - while cv2.getWindowProperty('frame', 0) >= 0: - ret, frame = cap.read() - assert ret, 'Video capture ended unexpectedly' + while True: + success, frame = cap.read() + if not success: + break frame = parser.annotate_frame(frame) + if FLAGS.video_path: + continue # Skip showing frame if parsing video. + cv2.imshow('frame', frame) keypress = cv2.waitKey(1) & 0xFF @@ -271,6 +281,9 @@ def main(argv): if keypress == ord('f'): parser.for_sale = not parser.for_sale + # Save any remaining unsaved items. + parser.save_items() + cap.release() cv2.destroyAllWindows()