class OfflineHextrack: def __init__(self, cfg, src, n, LED_pos, LED_tresholds): threading.current_thread().name = 'HexTrack' self.cfg = cfg self.frame_idx = 0 self.n = n self.mask_init = True self.made_mask = None # Create path to csv log file for tracking mouse position and LED-light state path = pkg_resources.resource_filename(__name__, "/data/interim/position_log_files/{}".format(src[len(src)-29:len(src)-10])) if not os.path.exists(path): try: os.mkdir(path) except OSError: print("Creation of the directory %s failed, this path probably already exists" % path) self.path = pkg_resources.resource_filename(__name__, '/data/interim/Position_log_files/{}/pos_log_file_{}.csv' .format(src[len(src)-29:len(src)-10], n)) # Initiation of the Grabbers and Trackers and creation of csv log file self.grabber = Grabber(src) self.tracker = Tracker(cfg, pos_log_file=open(self.path, 'w'), name=__name__, LED_pos=LED_pos, LED_thresholds=LED_tresholds) logging.debug('HexTrack initialization done!') self.vid = VideoFileClip(src) self.duration = self.vid.duration*15 # Loops through grabbing and tracking each frame of the video file def loop(self): pbar = tqdm(range(int(self.duration))) # pbar = tqdm(range(2000)) for i in pbar: frame = self.grabber.next() if frame is None: break # Checks if the frame has a mask already, if not, it creates a new mask if self.mask_init: self.tracker.apply(frame, self.frame_idx, n=self.n) elif not self.mask_init: self.tracker.apply(frame, self.frame_idx, mask_frame=self.made_mask, n=self.n) # At the second frame, show computer-generated mask # If not sufficient, gives possibility to input user-generated mask # if self.frame_idx == 1: # path = pkg_resources.resource_filename(__name__, '/output/Masks/mask_{}.png'.format(n)) # mask = cv2.imread(path) # plt.figure('Mask check') # plt.imshow(mask) # plt.show() # mask_check = input("If the mask is sufficient, enter y: ") # if mask_check != 'y': # input('Please upload custom mask under the name new_mask.png to the output folder and press enter') # self.made_mask = cv2.imread('new_mask.png', 0) # self.mask_init = False # self.frame_idx += 1 self.tracker.close() pbar.close() self.vid.reader.close() # Redundant, might be deleted later def process_events(self, display=False): if not display: return # Event loop call key = cv2.waitKey(1) # Process Keypress Events if key == ord('q'): self.stop() def stop(self): self.tracker.close() cv2.destroyAllWindows() raise SystemExit
class OfflineHextrack: def __init__(self, cfg, src, n, LED_pos, LED_thresholds, sources): # Video frame sources (top and bottom) self.sources = sources self.cfg = cfg self.frame_idx = 0 self.n = n self.mask_init = True self.made_mask = None # Create path to csv log file for tracking mouse position and LED-light state path = pkg_resources.resource_filename(__name__, "/data/interim/position_log_files/{}".format(src[len(src)-29: len(src)-10])) if not os.path.exists(path): try: os.mkdir(path) except OSError: print("Creation of the directory %s failed, this path probably already exists" % path) self.path = pkg_resources.resource_filename(__name__, '/data/interim/Position_log_files/{}/pos_log_file_{}.csv' .format(src[len(src)-29:len(src)-10], n)) # Initiation of the Grabbers and Trackers and creation of csv log file self.grabber = Grabber(src) self.tracker = Tracker(cfg, pos_log_file=open(self.path, 'w'), name=__name__, LED_pos=LED_pos, LED_thresholds=LED_thresholds) logging.debug('HexTrack initialization done!') # Video reader used to infer amount of frames self.vid = VideoFileClip(src) self.duration = self.vid.duration*15 self.src = src # Loops through grabbing and tracking each frame of the video file def loop(self): """"Loop through all frames in video and track mouse positions""" # tqdm package used to monitor tracking progress pbar = tqdm(range(int(self.duration))) for i in pbar: # Grab next frame, stops loop if no new frame is present (happens when all frames in video tracked) frame = self.grabber.next() if frame is None: break # Checks if the frame has a mask already, if not, it creates a new mask if self.mask_init: self.tracker.apply(frame, self.frame_idx, n=self.n, src=self.src) elif not self.mask_init: self.tracker.apply(frame, self.frame_idx, mask_frame=self.made_mask, n=self.n, src=self.src) if Mask_check: # At the second frame, show computer-generated mask # If not sufficient, gives possibility to input user-generated mask if self.frame_idx == 0: path = pkg_resources.resource_filename(__name__, "/data/raw/{}/Masks/mask_{}.png" .format(self.sources[0][len(self.sources[0])-29: len(self.sources[0])-10], n)) mask = cv2.imread(path) plt.figure('Mask check') plt.imshow(mask) plt.show() mask_check = input("If the mask is sufficient, enter y: ") if mask_check != 'y': input('Please upload custom mask under the name new_mask.png to the output folder' ' and press enter') mask_path = pkg_resources.resource_filename(__name__, "/Input_mask/new_mask.png") self.made_mask = cv2.imread(mask_path, 0) self.mask_init = False self.frame_idx += 1 # Close down tracker position log file, tqdm progress bar and video reader self.tracker.close() pbar.close() self.vid.reader.close() def stop(self): """Closes the position log files for following steps to be used""" self.tracker.close() cv2.destroyAllWindows() raise SystemExit