def export_tracks(self, database: TrackDatabase):
        """
        Writes tracks to a track database.
        :param database: database to write track to.
        """

        clip_id = os.path.basename(self.source_file)

        # overwrite any old clips.
        # Note: we do this even if there are no tracks so there there will be a blank clip entry as a record
        # that we have processed it.
        database.create_clip(clip_id, self)

        if len(self.tracks) == 0:
            return

        if not self.frame_buffer.has_flow:
            self.frame_buffer.generate_flow(self.opt_flow, self.flow_threshold)

        # get track data
        for track_number, track in enumerate(self.tracks):
            track_data = []
            for i in range(len(track)):
                channels = self.get_track_channels(track, i)

                # zero out the filtered channel
                if not self.INCLUDE_FILTERED_CHANNEL:
                    channels[TrackChannels.filtered] = 0
                track_data.append(channels)
            track_id = track_number+1
            database.add_track(clip_id, track_id, track_data, track, opts=blosc_zstd if self.ENABLE_COMPRESSION else None)
    def export_tracks(self, full_path, tracker: TrackExtractor, database: TrackDatabase):
        """
        Writes tracks to a track database.
        :param database: database to write track to.
        """

        clip_id = os.path.basename(full_path)

        # overwrite any old clips.
        # Note: we do this even if there are no tracks so there there will be a blank clip entry as a record
        # that we have processed it.
        database.create_clip(clip_id, tracker)

        if len(tracker.tracks) == 0:
            return

        tracker.generate_optical_flow()

        # get track data
        for track_number, track in enumerate(tracker.tracks):
            track_data = []
            for i in range(len(track)):
                channels = tracker.get_track_channels(track, i)

                # zero out the filtered channel
                if not self.config.extract.include_filtered_channel:
                    channels[TrackChannels.filtered] = 0
                track_data.append(channels)
            track_id = track_number+1
            start_time, end_time = tracker.start_and_end_time_absolute(track)
            database.add_track(clip_id, track_id, track_data,
                               track, opts=self.compression, start_time=start_time, end_time=end_time)
class CPTVTrackExtractor(CPTVFileProcessor):
    """
    Handles extracting tracks from CPTV files.
    Maintains a database recording which files have already been processed, and some statistics parameters used
    during processing.
    """

    def __init__(self, config, tracker_config):

        CPTVFileProcessor.__init__(self, config, tracker_config)

        self.hints = {}
        self.enable_track_output = True
        self.compression = (
            blosc_zstd if self.config.extract.enable_compression else None
        )

        self.previewer = Previewer.create_if_required(config, config.extract.preview)

        # normally poor quality tracks are filtered out, enabling this will let them through.
        self.disable_track_filters = False
        # disables background subtraction
        self.disable_background_subtraction = False

        os.makedirs(self.config.tracks_folder, mode=0o775, exist_ok=True)
        self.database = TrackDatabase(
            os.path.join(self.config.tracks_folder, "dataset.hdf5")
        )

        # load hints.  Hints are a way to give extra information to the tracker when necessary.
        # if os.path.exists(config.extract.hints_file):
        if config.extract.hints_file:
            self.load_hints(config.extract.hints_file)

    def load_hints(self, filename):
        """ Read in hints file from given path.  If file is not found an empty hints dictionary set."""

        self.hints = {}

        if not os.path.exists(filename):
            logging.warning("Failed to load hints file: %s", filename)
            return

        f = open(filename)
        for line_number, line in enumerate(f):
            line = line.strip()
            # comments
            if line == "" or line[0] == "#":
                continue
            try:
                (filename, file_max_tracks) = line.split()[:2]
            except:
                raise Exception("Error on line {0}: {1}".format(line_number, line))
            self.hints[filename] = int(file_max_tracks)

    def process_all(self, root):
        if root is None:
            root = self.config.source_folder

        previous_filter_setting = self.disable_track_filters
        previous_background_setting = self.disable_background_subtraction
        for folder_root, folders, _ in os.walk(root):

            for folder in folders:
                if folder not in self.config.excluded_folders:
                    if folder.lower() == "false-positive":
                        self.disable_track_filters = True
                        self.disable_background_subtraction = True
                        logging.info("Turning Track filters OFF.")

                    self.process_folder(
                        os.path.join(folder_root, folder), tag=folder.lower()
                    )

                    if folder.lower() == "false-positive":
                        logging.info("Restoring Track filters.")
                        self.disable_track_filters = previous_filter_setting
                        self.disable_background_subtraction = (
                            previous_background_setting
                        )

    def clean_tag(self, tag):
        """
        Removes all clips with given tag.
        :param tag: label to remove
        """
        logging.info("removing tag: %s", tag)

        ids = self.database.get_all_track_ids()
        for (clip_id, track_number) in ids:
            if not self.database.has_clip(clip_id):
                continue
            meta = self.database.get_track_meta(clip_id, track_number)
            if meta["tag"] == tag:
                logging.info("removing: %s", clip_id)
                self.database.remove_clip(clip_id)

    def clean_all(self):
        """
        Checks if there are any clips in the database that are on the banned list.  Also makes sure no track has more
        tracks than specified in hints file.
        """

        for clip_id, max_tracks in self.hints.items():
            if self.database.has_clip(clip_id):
                if max_tracks == 0:
                    logging.info(" - removing banned clip %s", clip_id)
                    self.database.remove_clip(clip_id)
                else:
                    meta = self.database.get_clip_meta(clip_id)
                    if meta["tracks"] > max_tracks:
                        logging.info(" - removing out of date clip: %s", clip_id)
                        self.database.remove_clip(clip_id)

    def process_file(self, full_path, **kwargs):
        """
        Extract tracks from specific file, and assign given tag.
        :param full_path: path: path to CPTV file to be processed
        :param tag: the tag to assign all tracks from this CPTV files
        :returns the tracker object
        """

        tag = kwargs["tag"]

        base_filename = os.path.splitext(os.path.split(full_path)[1])[0]
        cptv_filename = base_filename + ".cptv"

        logging.info(f"processing %s", cptv_filename)

        destination_folder = os.path.join(self.config.tracks_folder, tag.lower())
        os.makedirs(destination_folder, mode=0o775, exist_ok=True)
        # delete any previous files
        tools.purge(destination_folder, base_filename + "*.mp4")

        # read additional information from hints file
        if cptv_filename in self.hints:
            print(cptv_filename)
            logging.info(self.hints[cptv_filename])
            max_tracks = self.hints[cptv_filename]
            if max_tracks == 0:
                return
        else:
            max_tracks = self.config.tracking.max_tracks

        # load the track
        tracker = TrackExtractor(self.tracker_config)
        tracker.max_tracks = max_tracks
        tracker.tag = tag

        # by default we don't want to process the moving background images as it's too hard to get good tracks
        # without false-positives.
        tracker.reject_non_static_clips = True

        if self.disable_track_filters:
            tracker.track_min_delta = 0.0
            tracker.track_min_mass = 0.0
            tracker.track_min_offset = 0.0
            tracker.reject_non_static_clips = False

        if self.disable_background_subtraction:
            tracker.disable_background_subtraction = True

        # read metadata
        meta_data_filename = os.path.splitext(full_path)[0] + ".txt"
        if os.path.exists(meta_data_filename):

            meta_data = tools.load_clip_metadata(meta_data_filename)

            tags = set(
                [
                    tag["animal"]
                    for tag in meta_data["Tags"]
                    if "automatic" not in tag or not tag["automatic"]
                ]
            )

            # we can only handle one tagged animal at a time here.
            if len(tags) == 0:
                logging.warning(" - no tags in cptv files, ignoring.")
                return

            if len(tags) >= 2:
                # make sure all tags are the same
                logging.warning(" - mixed tags, can not process: %s", tags)
                return

            tracker.stats["confidence"] = meta_data["Tags"][0].get("confidence", 0.0)
            tracker.stats["trap"] = meta_data["Tags"][0].get("trap", "none")
            tracker.stats["event"] = meta_data["Tags"][0].get("event", "none")

            # clips tagged with false-positive sometimes come through with a null confidence rating
            # so we set it to 0.8 here.
            if (
                tracker.stats["event"] in ["false-positive", "false positive"]
                and tracker.stats["confidence"] is None
            ):
                tracker.stats["confidence"] = 0.8

            tracker.stats["cptv_metadata"] = meta_data
        else:
            self.log_warning(
                " - Warning: no tag metadata found for file - cannot use for machine learning."
            )

        start = time.time()

        # save some additional stats
        tracker.stats["version"] = TrackExtractor.VERSION

        tracker.load(full_path)

        if not tracker.extract_tracks():
            # this happens if the tracker rejected the video for some reason (i.e. too hot, or not static background).
            # we still need to make a record that we looked at it though.
            self.database.create_clip(os.path.basename(full_path), tracker)
            logging.warning(" - skipped (%s)", tracker.reject_reason)
            return tracker

        # assign each track the correct tag
        for track in tracker.tracks:
            track.tag = tag

        if self.enable_track_output:
            self.export_tracks(full_path, tracker, self.database)

        # write a preview
        if self.previewer:
            preview_filename = base_filename + "-preview" + ".mp4"
            preview_filename = os.path.join(destination_folder, preview_filename)
            self.previewer.create_individual_track_previews(preview_filename, tracker)
            self.previewer.export_clip_preview(preview_filename, tracker)

        if self.tracker_config.verbose:
            num_frames = len(tracker.frame_buffer.thermal)
            ms_per_frame = (time.time() - start) * 1000 / max(1, num_frames)
            self.log_message(
                "Tracks {}.  Frames: {}, Took {:.1f}ms per frame".format(
                    len(tracker.tracks), num_frames, ms_per_frame
                )
            )

        return tracker

    def export_tracks(
        self, full_path, tracks, tracker: TrackExtractor, database: TrackDatabase
    ):
        """
        Writes tracks to a track database.
        :param database: database to write track to.
        """

        clip_id = os.path.basename(full_path)

        # overwrite any old clips.
        # Note: we do this even if there are no tracks so there there will be a blank clip entry as a record
        # that we have processed it.
        database.create_clip(clip_id, tracker)

        if len(tracker.tracks) == 0:
            return

        tracker.generate_optical_flow()

        # get track data
        for track_number, track in enumerate(tracker.tracks):
            track_data = []
            for i in range(len(track)):
                channels = tracker.get_track_channels(track, i)

                # zero out the filtered channel
                if not self.config.extract.include_filtered_channel:
                    channels[TrackChannels.filtered] = 0
                track_data.append(channels)
            track_id = track_number + 1
            start_time, end_time = tracker.start_and_end_time_absolute(track)
            database.add_track(
                clip_id,
                track_id,
                track_data,
                track,
                opts=self.compression,
                start_time=start_time,
                end_time=end_time,
            )

    def needs_processing(self, source_filename):
        """
        Returns if given source file needs processing or not
        :param source_filename:
        :return:
        """

        clip_id = os.path.basename(source_filename)

        if self.config.reprocess:
            return True

        return not self.database.has_clip(clip_id)

    def run_test(self, source_folder, test: TrackerTestCase):
        """ Runs a specific test case. """

        def are_similar(value, expected, relative_error=0.2, abs_error=2.0):
            """ Checks of value is similar to expected value. An expected value of 0 will always return true. """
            if expected == 0:
                return True
            return ((abs(value - expected) / expected) <= relative_error) or (
                abs(value - expected) <= abs_error
            )

        # find the file.  We looking in all the tag folder to make life simpler when creating the test file.
        source_file = tools.find_file(source_folder, test.source)

        # make sure we don't write to database
        self.enable_track_output = False

        if source_file is None:
            logging.warning(
                "Could not find %s in root folder %s", test.source, source_folder
            )
            return

        logging.info(source_file)
        tracker = self.process_file(source_file, tag="test")

        # read in stats files and see how we did
        if len(tracker.tracks) != len(test.tracks):
            logging.error(
                "%s Incorrect number of tracks, expected %s found %s",
                test.source,
                len(test.tracks),
                len(tracker.tracks),
            )
            return

        for test_result, (expected_duration, expected_movement) in zip(
            tracker.tracks, test.tracks
        ):

            track_stats = test_result.get_stats()

            if not are_similar(
                test_result.duration, expected_duration
            ) or not are_similar(track_stats.max_offset, expected_movement):
                logging.error(
                    "%s Track too dissimilar expected %s but found %s",
                    test.source,
                    (expected_duration, expected_movement),
                    (test_result.duration, track_stats.max_offset),
                )
            else:
                logging.info("%s passed", test.source)

    def run_tests(self, source_folder, tests_file):
        """ Processes file in test file and compares results to expected output. """

        # disable hints for tests
        self.hints = []

        tests = []
        test = None

        # # we need to make sure all tests are redone every time.
        # self.overwrite_mode = self.OM_ALL

        # load in the test data
        for line in open(tests_file, "r"):
            line = line.strip()
            if line == "":
                continue
            if line[0] == "#":
                continue

            if line.split()[0].lower() == "track":
                if test == None:
                    raise Exception("Can not have track before source file.")
                expected_length, expected_movement = [int(x) for x in line.split()[1:]]
                test.tracks.append((expected_length, expected_movement))
            else:
                test = TrackerTestCase()
                test.source = line
                tests.append(test)

        logging.info("Found %d test cases", len(tests))

        for test in tests:
            self.run_test(source_folder, test)
class CPTVTrackExtractor(CPTVFileProcessor):
    """
    Handles extracting tracks from CPTV files.
    Maintains a database recording which files have already been processed, and some statistics parameters used
    during processing.
    """

    # version number.  Recorded into stats file when a clip is processed.
    VERSION = 6

    def __init__(self, out_folder):

        CPTVFileProcessor.__init__(self)

        self.hints = {}
        self.colormap = plt.get_cmap('jet')
        self.verbose = False
        self.out_folder = out_folder
        self.overwrite_mode = CPTVTrackExtractor.OM_NONE
        self.enable_previews = False
        self.enable_track_output = True

        # normally poor quality tracks are filtered out, enabling this will let them through.
        self.disable_track_filters = False
        # disables background subtraction
        self.disable_background_subtraction = False

        self.high_quality_optical_flow = False

        self.database = TrackDatabase(os.path.join(self.out_folder, 'dataset.hdf5'))

        self.worker_pool_init = init_workers

    def load_hints(self, filename):
        """ Read in hints file from given path.  If file is not found an empty hints dictionary set."""

        self.hints = {}

        if not os.path.exists(filename):
            return

        f = open(filename)
        for line_number, line in enumerate(f):
            line = line.strip()
            # comments
            if line == '' or line[0] == '#':
                continue
            try:
                (filename, file_max_tracks) = line.split()[:2]
            except:
                raise Exception("Error on line {0}: {1}".format(line_number, line))
            self.hints[filename] = int(file_max_tracks)

    def process_all(self, root):

        previous_filter_setting = self.disable_track_filters
        previous_background_setting = self.disable_background_subtraction

        for root, folders, files in os.walk(root):
            for folder in folders:
                if folder not in EXCLUDED_FOLDERS:
                    if folder.lower() == "false-positive":
                        self.disable_track_filters = True
                        self.disable_background_subtraction = True
                        print("Turning Track filters OFF.")
                    self.process_folder(os.path.join(root, folder), tag=folder.lower(), worker_pool_args=(trackdatabase.hdf5_lock,))
                    if folder.lower() == "false-positive":
                        print("Restoring Track filters.")
                        self.disable_track_filters = previous_filter_setting
                        self.disable_background_subtraction = previous_background_setting



    def clean_tag(self, tag):
        """
        Removes all clips with given tag.
        :param tag: label to remove
        """
        print("removing tag {}".format(tag))

        ids = self.database.get_all_track_ids()
        for (clip_id, track_number) in ids:
            if not self.database.has_clip(clip_id):
                continue
            meta = self.database.get_track_meta(clip_id, track_number)
            if meta['tag'] == tag:
                print("removing", clip_id)
                self.database.remove_clip(clip_id)


    def clean_all(self):
        """
        Checks if there are any clips in the database that are on the banned list.  Also makes sure no track has more
        tracks than specified in hints file.
        """

        for clip_id, max_tracks in self.hints.items( ):
            if self.database.has_clip(clip_id):
                if max_tracks == 0:
                    print(" - removing banned clip {}".format(clip_id))
                    self.database.remove_clip(clip_id)
                else:
                    meta = self.database.get_clip_meta(clip_id)
                    if meta['tracks'] > max_tracks:
                        print(" - removing out of date clip {}".format(clip_id))
                        self.database.remove_clip(clip_id)


    def process_file(self, full_path, **kwargs):
        """
        Extract tracks from specific file, and assign given tag.
        :param full_path: path: path to CPTV file to be processed
        :param tag: the tag to assign all tracks from this CPTV files
        :param create_preview_file: if enabled creates an MPEG preview file showing the tracking working.  This
            process can be quite time consuming.
        :returns the tracker object
        """

        tag = kwargs['tag']

        base_filename = os.path.splitext(os.path.split(full_path)[1])[0]
        cptv_filename = base_filename + '.cptv'
        preview_filename = base_filename + '-preview' + '.mp4'
        stats_filename = base_filename + '.txt'

        destination_folder = os.path.join(self.out_folder, tag.lower())

        stats_path_and_filename = os.path.join(destination_folder, stats_filename)

        # read additional information from hints file
        if cptv_filename in self.hints:
            max_tracks = self.hints[cptv_filename]
            if max_tracks == 0:
                return
        else:
            max_tracks = 10

        # make destination folder if required
        try:
            os.stat(destination_folder)
        except:
            self.log_message(" Making path " + destination_folder)
            os.mkdir(destination_folder)

        # check if we have already processed this file
        if self.needs_processing(stats_path_and_filename):
            print("Processing {0} [{1}]".format(cptv_filename, tag))
        else:
            return

        # delete any previous files
        tools.purge(destination_folder, base_filename + "*.mp4")

        # load the track
        tracker = TrackExtractor()
        tracker.max_tracks = max_tracks
        tracker.tag = tag
        tracker.verbose = self.verbose >= 2
        tracker.high_quality_optical_flow = self.high_quality_optical_flow

        # by default we don't want to process the moving background images as it's too hard to get good tracks
        # without false-positives.
        tracker.reject_non_static_clips = True

        if self.disable_track_filters:
            tracker.track_min_delta = 0.0
            tracker.track_min_mass = 0.0
            tracker.track_min_offset = 0.0
            tracker.reject_non_static_clips = False

        if self.disable_background_subtraction:
            tracker.disable_background_subtraction = True

        # read metadata
        meta_data_filename = os.path.splitext(full_path)[0] + ".txt"
        if os.path.exists(meta_data_filename):

            meta_data = tools.load_clip_metadata(meta_data_filename)

            tags = set([tag['animal'] for tag in meta_data['Tags'] if 'automatic' not in tag or not tag['automatic']])

            # we can only handle one tagged animal at a time here.
            if len(tags) == 0:
                print(" - Warning, no tags in cptv files, ignoring.")
                return

            if len(tags)>= 2:
                # make sure all tags are the same
                print(" - Warning, mixed tags, can not process.",tags)
                return

            tracker.stats['confidence'] = meta_data['Tags'][0].get('confidence',0.0)
            tracker.stats['trap'] = meta_data['Tags'][0].get('trap','none')
            tracker.stats['event'] = meta_data['Tags'][0].get('event','none')

            # clips tagged with false-positive sometimes come through with a null confidence rating
            # so we set it to 0.8 here.
            if tracker.stats['event'] in ['false-positive', 'false positive'] and tracker.stats['confidence'] is None:
                tracker.stats['confidence'] = 0.8

            tracker.stats['cptv_metadata'] = meta_data
        else:
            self.log_warning(" - Warning: no metadata found for file.")
            return

        start = time.time()

        # save some additional stats
        tracker.stats['version'] = CPTVTrackExtractor.VERSION

        tracker.load(full_path)

        if not tracker.extract_tracks():
            # this happens if the tracker rejected the video for some reason (i.e. too hot, or not static background).
            # we still need to make a record that we looked at it though.
            self.database.create_clip(os.path.basename(full_path), tracker)
            print(" - skipped ({})".format(tracker.reject_reason))
            return tracker

        # assign each track the correct tag
        for track in tracker.tracks:
            track.tag = tag

        if self.enable_track_output:
            tracker.export_tracks(self.database)

        # write a preview
        if self.enable_previews:
            self.export_mpeg_preview(os.path.join(destination_folder, preview_filename), tracker)

        time_per_frame = (time.time() - start) / len(tracker.frame_buffer)

        # time_stats = tracker.stats['time_per_frame']
        self.log_message(" -tracks: {} {:.1f}sec - Time per frame: {:.1f}ms".format(
             len(tracker.tracks),
             sum(track.duration for track in tracker.tracks),
             time_per_frame * 1000
         ))

        return tracker

    def needs_processing(self, source_filename):
        """
        Returns if given source file needs processing or not
        :param source_filename:
        :return:
        """

        clip_id = os.path.basename(source_filename)

        if self.overwrite_mode == self.OM_ALL:
            return True

        return not self.database.has_clip(clip_id)

    def run_test(self, source_folder, test: TrackerTestCase):
        """ Runs a specific test case. """

        def are_similar(value, expected, relative_error = 0.2, abs_error = 2.0):
            """ Checks of value is similar to expected value. An expected value of 0 will always return true. """
            if expected == 0:
                return True
            return ((abs(value - expected) / expected) <= relative_error) or (abs(value - expected) <= abs_error)

        # find the file.  We looking in all the tag folder to make life simpler when creating the test file.
        source_file = tools.find_file(source_folder, test.source)

        # make sure we don't write to database
        self.enable_track_output = False

        if source_file is None:
            print("Could not find {0} in root folder {1}".format(test.source, source_folder))
            return

        print(source_file)
        tracker = self.process_file(source_file, tag='test')

        # read in stats files and see how we did
        if len(tracker.tracks) != len(test.tracks):
            print("[Fail] {0} Incorrect number of tracks, expected {1} found {2}".format(test.source, len(test.tracks), len(tracker.tracks)))
            return

        for test_result, (expected_duration, expected_movement) in zip(tracker.tracks, test.tracks):

            track_stats = test_result.get_stats()

            if not are_similar(test_result.duration, expected_duration) or not are_similar(track_stats.max_offset, expected_movement):
                print("[Fail] {0} Track too dissimilar expected {1} but found {2}".format(
                    test.source,
                    (expected_duration, expected_movement),
                    (test_result.duration, track_stats.max_offset)))
            else:
                print("[PASS] {0}".format(test.source))

    def export_track_mpeg_previews(self, filename_base, tracker: TrackExtractor):
        """
        Exports preview MPEG for a specific track
        :param filename_base:
        :param tracker:
        :param track:
        :return:
        """

        # resolution of video file.
        # videos look much better scaled up
        FRAME_SIZE = 4*48

        frame_width, frame_height = FRAME_SIZE, FRAME_SIZE
        frame_width =  frame_width // 4 * 4
        frame_height = frame_height // 4 * 4

        for id, track in enumerate(tracker.tracks):
            video_frames = []
            for frame_number in range(len(track.bounds_history)):
                channels = tracker.get_track_channels(track, frame_number)
                img = tools.convert_heat_to_img(channels[1], self.colormap, 0, 350)
                img = img.resize((frame_width, frame_height), Image.NEAREST)
                video_frames.append(np.asarray(img))

            tools.write_mpeg(filename_base+"-"+str(id+1)+".mp4", video_frames)

    def export_mpeg_preview(self, filename, tracker: TrackExtractor):
        """
        Exports tracking information preview to MPEG file.
        """

        self.export_track_mpeg_previews(os.path.splitext(filename)[0], tracker)

        MPEGStreamer = MPEGPreviewStreamer(tracker, self.colormap)

        tools.stream_mpeg(filename, MPEGStreamer)

    def run_tests(self, source_folder, tests_file):
        """ Processes file in test file and compares results to expected output. """

        # disable hints for tests
        self.hints = []

        tests = []
        test = None

        # we need to make sure all tests are redone every time.
        self.overwrite_mode = CPTVTrackExtractor.OM_ALL

        # load in the test data
        for line in open(tests_file, 'r'):
            line = line.strip()
            if line == '':
                continue
            if line[0] == '#':
                continue

            if line.split()[0].lower() == 'track':
                if test == None:
                    raise Exception("Can not have track before source file.")
                expected_length, expected_movement = [int(x) for x in line.split()[1:]]
                test.tracks.append((expected_length, expected_movement))
            else:
                test = TrackerTestCase()
                test.source = line
                tests.append(test)

        print("Found {0} test cases".format(len(tests)))

        for test in tests:
            self.run_test(source_folder, test)
class ClipLoader:
    def __init__(self, config):

        self.config = config
        os.makedirs(self.config.tracks_folder, mode=0o775, exist_ok=True)
        self.database = TrackDatabase(
            os.path.join(self.config.tracks_folder, "dataset.hdf5"))

        self.compression = (tools.gzip_compression
                            if self.config.load.enable_compression else None)
        self.track_config = config.tracking
        # number of threads to use when processing jobs.
        self.workers_threads = config.worker_threads
        self.previewer = Previewer.create_if_required(config,
                                                      config.load.preview)
        self.track_extractor = ClipTrackExtractor(
            self.config.tracking,
            self.config.use_opt_flow
            or config.load.preview == Previewer.PREVIEW_TRACKING,
            self.config.load.cache_to_disk,
        )

    def process_all(self, root=None):
        if root is None:
            root = self.config.source_folder

        jobs = []
        for folder_path, _, files in os.walk(root):
            for name in files:
                if os.path.splitext(name)[1] == ".cptv":
                    full_path = os.path.join(folder_path, name)
                    jobs.append((self, full_path))

        self._process_jobs(jobs)

    def _process_jobs(self, jobs):
        if self.workers_threads == 0:
            for job in jobs:
                process_job(job)
        else:
            pool = multiprocessing.Pool(self.workers_threads)
            try:
                pool.map(process_job, jobs, chunksize=1)
                pool.close()
                pool.join()
            except KeyboardInterrupt:
                logging.info("KeyboardInterrupt, terminating.")
                pool.terminate()
                exit()
            except Exception:
                logging.exception("Error processing files")
            else:
                pool.close()

    def _get_dest_folder(self, filename):
        return os.path.join(self.config.tracks_folder,
                            get_distributed_folder(filename))

    def _export_tracks(self, full_path, clip):
        """
        Writes tracks to a track database.
        :param database: database to write track to.
        """
        # overwrite any old clips.
        # Note: we do this even if there are no tracks so there there will be a blank clip entry as a record
        # that we have processed it.
        self.database.create_clip(clip)

        for track in clip.tracks:
            start_time, end_time = clip.start_and_end_time_absolute(
                track.start_s, track.end_s)
            track_data = []
            for region in track.bounds_history:
                frame = clip.frame_buffer.get_frame(region.frame_number)
                frame = track.crop_by_region(frame, region)
                # zero out the filtered channel
                if not self.config.load.include_filtered_channel:
                    frame[TrackChannels.filtered] = 0
                track_data.append(frame)

            self.database.add_track(
                clip.get_id(),
                track,
                track_data,
                opts=self.compression,
                start_time=start_time,
                end_time=end_time,
            )

    def _filter_clip_tracks(self, clip_metadata):
        """
        Removes track metadata for tracks which are invalid. Tracks are invalid
        if they aren't confident or they are in the excluded_tags list.
        Returns valid tracks
        """

        tracks_meta = clip_metadata.get("Tracks", [])
        valid_tracks = [
            track for track in tracks_meta if self._track_meta_is_valid(track)
        ]
        clip_metadata["Tracks"] = valid_tracks
        return valid_tracks

    def _track_meta_is_valid(self, track_meta):
        """ 
        Tracks are valid if their confidence meets the threshold and they are
        not in the excluded_tags list, defined in the config.
        """
        min_confidence = self.track_config.min_tag_confidence
        excluded_tags = self.config.excluded_tags
        track_data = track_meta.get("data")
        if not track_data:
            return False

        track_tag = Track.get_best_human_tag(track_meta,
                                             self.config.load.tag_precedence,
                                             min_confidence)
        if track_tag is None:
            return False
        tag = track_tag.get("what")
        confidence = track_tag.get("confidence", 0)
        return tag and tag not in excluded_tags and confidence >= min_confidence

    def process_file(self, filename):
        start = time.time()
        base_filename = os.path.splitext(os.path.basename(filename))[0]

        logging.info(f"processing %s", filename)

        destination_folder = self._get_dest_folder(base_filename)
        os.makedirs(destination_folder, mode=0o775, exist_ok=True)
        # delete any previous files
        tools.purge(destination_folder, base_filename + "*.mp4")

        # read metadata
        metadata_filename = os.path.join(os.path.dirname(filename),
                                         base_filename + ".txt")

        if not os.path.isfile(metadata_filename):
            logging.error("No meta data found for %s", metadata_filename)
            return

        metadata = tools.load_clip_metadata(metadata_filename)
        valid_tracks = self._filter_clip_tracks(metadata)
        if not valid_tracks:
            logging.error("No valid track data found for %s", filename)
            return

        clip = Clip(self.track_config, filename)
        clip.load_metadata(
            metadata,
            self.config.load.include_filtered_channel,
            self.config.load.tag_precedence,
        )

        self.track_extractor.parse_clip(clip)
        # , self.config.load.cache_to_disk, self.config.use_opt_flow

        if self.track_config.enable_track_output:
            self._export_tracks(filename, clip)

        # write a preview
        if self.previewer:
            preview_filename = base_filename + "-preview" + ".mp4"
            preview_filename = os.path.join(destination_folder,
                                            preview_filename)
            self.previewer.create_individual_track_previews(
                preview_filename, clip)
            self.previewer.export_clip_preview(preview_filename, clip)

        if self.track_config.verbose:
            num_frames = len(clip.frame_buffer.frames)
            ms_per_frame = (time.time() - start) * 1000 / max(1, num_frames)
            self._log_message(
                "Tracks {}.  Frames: {}, Took {:.1f}ms per frame".format(
                    len(clip.tracks), num_frames, ms_per_frame))

    def _log_message(self, message):
        """ Record message in stdout.  Will be printed if verbose is enabled. """
        # note, python has really good logging... I should probably make use of this.
        if self.track_config.verbose:
            logging.info(message)
Example #6
0
class ClipLoader:
    def __init__(self, config, reprocess=False):

        self.config = config
        os.makedirs(self.config.tracks_folder, mode=0o775, exist_ok=True)
        self.database = TrackDatabase(
            os.path.join(self.config.tracks_folder, "dataset.hdf5")
        )
        self.reprocess = reprocess
        self.compression = (
            tools.gzip_compression if self.config.load.enable_compression else None
        )
        self.track_config = config.tracking
        # number of threads to use when processing jobs.
        self.workers_threads = config.worker_threads
        self.previewer = Previewer.create_if_required(config, config.load.preview)
        self.track_extractor = ClipTrackExtractor(
            self.config.tracking,
            self.config.use_opt_flow
            or config.load.preview == Previewer.PREVIEW_TRACKING,
            self.config.load.cache_to_disk,
            high_quality_optical_flow=self.config.load.high_quality_optical_flow,
            verbose=config.verbose,
        )

    def process_all(self, root):
        job_queue = Queue()
        processes = []
        for i in range(max(1, self.workers_threads)):
            p = Process(
                target=process_job,
                args=(self, job_queue),
            )
            processes.append(p)
            p.start()
        if root is None:
            root = self.config.source_folder
        file_paths = []
        for folder_path, _, files in os.walk(root):
            for name in files:
                if os.path.splitext(name)[1] == ".cptv":
                    full_path = os.path.join(folder_path, name)
                    file_paths.append(full_path)
        # allows us know the order of processing
        file_paths.sort()
        for file_path in file_paths:
            job_queue.put(file_path)

        logging.info("Processing %d", job_queue.qsize())
        for i in range(len(processes)):
            job_queue.put("DONE")
        for process in processes:
            try:
                process.join()
            except KeyboardInterrupt:
                logging.info("KeyboardInterrupt, terminating.")
                for process in processes:
                    process.terminate()
                exit()

    def _get_dest_folder(self, filename):
        return os.path.join(self.config.tracks_folder, get_distributed_folder(filename))

    def _export_tracks(self, full_path, clip):
        """
        Writes tracks to a track database.
        :param database: database to write track to.
        """
        # overwrite any old clips.
        # Note: we do this even if there are no tracks so there there will be a blank clip entry as a record
        # that we have processed it.
        self.database.create_clip(clip)
        for track in clip.tracks:

            start_time, end_time = clip.start_and_end_time_absolute(
                track.start_s, track.end_s
            )
            original_thermal = []
            cropped_data = []
            for region in track.bounds_history:
                frame = clip.frame_buffer.get_frame(region.frame_number)
                original_thermal.append(frame.thermal)
                cropped = frame.crop_by_region(region)
                # zero out the filtered channel
                if not self.config.load.include_filtered_channel:
                    cropped.filtered = np.zeros(cropped.thermal.shape)
                cropped_data.append(cropped)

            sample_frames = get_sample_frames(
                clip.ffc_frames,
                [bounds.mass for bounds in track.bounds_history],
                self.config.build.segment_min_avg_mass,
                cropped_data,
            )
            self.database.add_track(
                clip.get_id(),
                track,
                cropped_data,
                sample_frames=sample_frames,
                opts=self.compression,
                original_thermal=original_thermal,
                start_time=start_time,
                end_time=end_time,
            )
        self.database.finished_processing(clip.get_id())

    def _filter_clip_tracks(self, clip_metadata):
        """
        Removes track metadata for tracks which are invalid. Tracks are invalid
        if they aren't confident or they are in the excluded_tags list.
        Returns valid tracks
        """

        tracks_meta = clip_metadata.get("Tracks", [])
        valid_tracks = [
            track for track in tracks_meta if self._track_meta_is_valid(track)
        ]
        clip_metadata["Tracks"] = valid_tracks
        return valid_tracks

    def _track_meta_is_valid(self, track_meta):
        """
        Tracks are valid if their confidence meets the threshold and they are
        not in the excluded_tags list, defined in the config.
        """
        min_confidence = self.track_config.min_tag_confidence
        track_data = track_meta.get("data")
        if not track_data:
            return False

        track_tags = []
        if "TrackTags" in track_meta:
            track_tags = track_meta["TrackTags"]
        excluded_tags = [
            tag
            for tag in track_tags
            if not tag.get("automatic", False) and tag in self.config.load.excluded_tags
        ]

        if len(excluded_tags) > 0:
            return False

        track_tag = Track.get_best_human_tag(
            track_tags, self.config.load.tag_precedence, min_confidence
        )
        if track_tag is None:
            return False
        tag = track_tag.get("what")
        confidence = track_tag.get("confidence", 0)
        return tag and tag not in excluded_tags and confidence >= min_confidence

    def process_file(self, filename, classifier=None):
        start = time.time()
        base_filename = os.path.splitext(os.path.basename(filename))[0]

        logging.info(f"processing %s", filename)

        destination_folder = self._get_dest_folder(base_filename)
        # delete any previous files
        tools.purge(destination_folder, base_filename + "*.mp4")

        # read metadata
        metadata_filename = os.path.join(
            os.path.dirname(filename), base_filename + ".txt"
        )

        if not os.path.isfile(metadata_filename):
            logging.error("No meta data found for %s", metadata_filename)
            return

        metadata = tools.load_clip_metadata(metadata_filename)
        if not self.reprocess and self.database.has_clip(str(metadata["id"])):
            logging.warning("Already loaded %s", filename)
            return

        valid_tracks = self._filter_clip_tracks(metadata)
        if not valid_tracks or len(valid_tracks) == 0:
            logging.error("No valid track data found for %s", filename)
            return

        clip = Clip(self.track_config, filename)
        clip.load_metadata(
            metadata,
            self.config.load.tag_precedence,
        )
        tracker_versions = set(
            [
                t.get("data", {}).get("tracker_version", 0)
                for t in metadata.get("Tracks", [])
            ]
        )
        if len(tracker_versions) > 1:
            logginer.error(
                "Tracks with different tracking versions cannot process %s versions %s",
                filename,
                tracker_versions,
            )
            return
        tracker_version = tracker_versions.pop()
        process_background = tracker_version < 10
        logging.debug(
            "Processing background? %s tracker_versions %s",
            process_background,
            tracker_version,
        )
        if not self.track_extractor.parse_clip(
            clip, process_background=process_background
        ):
            logging.error("No valid clip found for %s", filename)
            return

        # , self.config.load.cache_to_disk, self.config.use_opt_flow
        if self.track_config.enable_track_output:
            self._export_tracks(filename, clip)

        # write a preview
        if self.previewer:
            os.makedirs(destination_folder, mode=0o775, exist_ok=True)

            preview_filename = base_filename + "-preview" + ".mp4"
            preview_filename = os.path.join(destination_folder, preview_filename)
            self.previewer.create_individual_track_previews(preview_filename, clip)
            self.previewer.export_clip_preview(preview_filename, clip)

        if self.config.verbose:
            num_frames = len(clip.frame_buffer.frames)
            ms_per_frame = (time.time() - start) * 1000 / max(1, num_frames)
            self._log_message(
                "Tracks {}.  Frames: {}, Took {:.1f}ms per frame".format(
                    len(clip.tracks), num_frames, ms_per_frame
                )
            )

    def _log_message(self, message):
        """Record message in stdout.  Will be printed if verbose is enabled."""
        # note, python has really good logging... I should probably make use of this.
        if self.track_config.verbose:
            logging.info(message)