OTVision 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

OTVision/track/track.py

@@ -1,29 +1,5 @@
-"""
-OTVision main module for tracking objects in successive frames of videos
-"""
-
-# points and transform tracksectory points from pixel into world coordinates.
-
 import logging
 import uuid
-
-# Copyright (C) 2022 OpenTrafficCam Contributors
-# <https://github.com/OpenTrafficCam
-# <team@opentrafficcam.org>
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more detectionsails.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <https://www.gnu.org/licenses/>.
-from collections import defaultdict
 from pathlib import Path
 from typing import Callable, Iterator
 
@@ -32,7 +8,6 @@ from tqdm import tqdm
 from OTVision import dataformat
 from OTVision.config import (
     CONFIG,
-    DEFAULT_FILETYPE,
     DETECT,
     FILETYPES,
     IOU,
@@ -44,103 +19,43 @@ from OTVision.config import (
     T_MISS_MAX,
     TRACK,
 )
-from OTVision.dataformat import DATA, DETECTIONS, FINISHED, METADATA, TRACK_ID
-from OTVision.helpers.files import denormalize_bbox, get_files, write_json
+from OTVision.helpers.files import get_files
+from OTVision.helpers.input_types import check_types
 from OTVision.helpers.log import LOGGER_NAME
-from OTVision.track.preprocess import (
-    FrameChunk,
-    FrameChunkParser,
-    FrameIndexer,
-    Preprocess,
+from OTVision.track.exporter.filebased_exporter import FinishedChunkTrackExporter
+from OTVision.track.parser.chunk_parser_plugins import JsonChunkParser
+from OTVision.track.parser.frame_group_parser_plugins import (
+    TimeThresholdFrameGroupParser,
 )
+from OTVision.track.tracker.filebased_tracking import (
+    GroupedFilesTracker,
+    UnfinishedChunksBuffer,
+)
+from OTVision.track.tracker.tracker_plugin_iou import IouParameters, IouTracker
 
-from .iou import TrackedDetections, TrackingResult, id_generator, track_iou
-
-log = logging.getLogger(LOGGER_NAME)
-
-IdGenerator = Callable[[], str]
-
-
-class TrackedChunk(TrackedDetections):
-    """Tracking results combined with respective metadata."""
-
-    def __init__(
-        self,
-        detections: TrackedDetections,
-        file_path: Path,
-        frame_group_id: int,
-        frame_offset: int,
-        metadata: dict,
-    ) -> None:
-        super().__init__(
-            detections._detections,
-            detections._detected_ids,
-            detections._active_track_ids,
-        )
-        self.file_path = file_path
-        self.frame_offset = frame_offset
-        self._frame_group_id = frame_group_id
-        self.metadata = metadata
-
-
-class TrackingResultStore:
-    """TrackingResultStore manages TrackedChunks
-    and data required for the next tracking iteration:
-    - the remaining active tracks of the last iteration
-    - a global lookup table for each tracks last detection frame
-    """
-
-    def __init__(self, tracking_run_id: str) -> None:
-        self.tracking_run_id = tracking_run_id
-        self._tracked_chunks: list[TrackedChunk] = list()
-        self.last_track_frame: dict[int, int] = dict()
-        self.active_tracks: list[dict] = []
-
-    def store_tracking_result(
-        self,
-        tracking_result: TrackingResult,
-        file_path: Path,
-        frame_group_id: int,
-        frame_offset: int,
-        metadata: dict,
-    ) -> None:
-        # store tracking results
-        self._tracked_chunks.append(
-            TrackedChunk(
-                detections=tracking_result.tracked_detections,
-                file_path=file_path,
-                frame_group_id=frame_group_id,
-                frame_offset=frame_offset,
-                metadata=metadata,
-            )
-        )
-        # update global lookup table
-        self.last_track_frame.update(tracking_result.last_track_frame)
-
-        # replace last iterations active tracks
-        self.active_tracks = tracking_result.active_tracks
-        print("remaining active tracks", len(self.active_tracks))
 
-        # collect remaining active ids and update tracked chunks
-        new_active_ids = {t[TRACK_ID] for t in self.active_tracks}
-        for result in self._tracked_chunks:
-            result.update_active_track_ids(new_active_ids)
+def track_id_generator() -> Iterator[int]:
+    ID: int = 0
+    while True:
+        ID += 1
+        yield ID
 
-    def get_finished_results(self) -> list[TrackedChunk]:
-        """Return all TrackedChunks that have no remaining active track ids.
-        These are ready to be written to a result file.
 
-        Returns:
-            list[TrackedChunk]: finished TrackChunks
-        """
-        finished = [chunk for chunk in self._tracked_chunks if chunk.is_finished()]
+log = logging.getLogger(LOGGER_NAME)
+STR_ID_GENERATOR = Callable[[], str]
 
-        # remove finished TrackedChunks from manager
-        self._tracked_chunks = [
-            chunk for chunk in self._tracked_chunks if chunk not in finished
-        ]
 
-        return finished
+def tracker_metadata(
+    sigma_l: float, sigma_h: float, sigma_iou: float, t_min: float, t_miss_max: float
+) -> dict:
+    return {
+        dataformat.NAME: "IOU",
+        dataformat.SIGMA_L: sigma_l,
+        dataformat.SIGMA_H: sigma_h,
+        dataformat.SIGMA_IOU: sigma_iou,
+        dataformat.T_MIN: t_min,
+        dataformat.T_MISS_MAX: t_miss_max,
+    }
 
 
 def main(
@@ -151,7 +66,7 @@ def main(
     t_min: int = CONFIG[TRACK][IOU][T_MIN],
     t_miss_max: int = CONFIG[TRACK][IOU][T_MISS_MAX],
     overwrite: bool = CONFIG[TRACK][OVERWRITE],
-    tracking_run_id_generator: IdGenerator = lambda: str(uuid.uuid4()),
+    tracking_run_id_generator: STR_ID_GENERATOR = lambda: str(uuid.uuid4()),
 ) -> None:
     """Read detections from otdet file, perform tracking using iou tracker and
         save tracks to ottrk file.
@@ -183,6 +98,8 @@ def main(
             for this tracking run
     """
 
+    check_types(sigma_l, sigma_h, sigma_iou, t_min, t_miss_max)
+
     filetypes = CONFIG[FILETYPES][DETECT]
     detections_files = get_files(paths=paths, filetypes=filetypes)
 
@@ -194,229 +111,34 @@ def main(
         log.warning(f"No files of type '{filetypes}' found to track!")
         return
 
-    tracking_run_id = tracking_run_id_generator()
-    preprocessor = Preprocess()
-    preprocessed = preprocessor.run(detections_files)
-    file_type = CONFIG[DEFAULT_FILETYPE][TRACK]
-
-    for frame_group_id, frame_group in tqdm(
-        enumerate(preprocessed),
-        desc="Tracked frame groups",
-        unit=" framegroup",
-    ):
-        print()
-        print(f"Process frame group {frame_group_id}")
-
-        # update metadata for all files in FrameGroup
-        frame_group.update_metadata(
-            tracker_metadata(sigma_l, sigma_h, sigma_iou, t_min, t_miss_max)
-        )
-
-        # metadata and tracking results to keep track during iterations
-        # within a FrameGroups track/vehicle and frame ids should be unique
-        vehicle_id_generator = id_generator()
-        frame_offset = 0
-        track_result_store = TrackingResultStore(tracking_run_id)
-
-        # process each otdet file in frame group
-        for file_path in frame_group.files:
-            print(f"Process file {file_path} in frame group {frame_group_id}")
-
-            # read detection data
-            chunk = FrameChunkParser.parse(file_path, frame_offset)
-
-            if skip_existing_output_files(chunk, overwrite, file_type):
-                continue
-
-            log.info(f"Track {str(chunk)}")
-
-            detections = chunk.to_dict()
-            detections_denormalized = denormalize_bbox(
-                detections, metadata=frame_group._files_metadata
-            )
-
-            tracking_result = track(
-                detections=detections_denormalized,
-                sigma_l=sigma_l,
-                sigma_h=sigma_h,
-                sigma_iou=sigma_iou,
-                t_min=t_min,
-                t_miss_max=t_miss_max,
-                previous_active_tracks=track_result_store.active_tracks,
-                vehicle_id_generator=vehicle_id_generator,
-            )
-
-            track_result_store.store_tracking_result(
-                tracking_result,
-                file_path=file_path,
-                frame_group_id=frame_group_id,
-                frame_offset=frame_offset,
-                metadata=frame_group.metadata_for(file_path),
-            )
-            log.debug(f"Successfully tracked {chunk}")
-            frame_offset = chunk.last_frame_id() + 1
-
-            for finished_chunk in track_result_store.get_finished_results():
-                mark_and_write_results(finished_chunk, track_result_store, overwrite)
-
-        # write last files of frame group
-        # even though some tracks mights still be active
-        for finished_chunk in track_result_store._tracked_chunks:
-            mark_and_write_results(finished_chunk, track_result_store, overwrite)
-
-        log.info("Successfully tracked and wrote ")
-
-    finished_msg = "Finished tracking"
-    log.info(finished_msg)
-    print(finished_msg)
-
-
-def skip_existing_output_files(
-    chunk: FrameChunk, overwrite: bool, file_type: str
-) -> bool:
-    existing_output_files = chunk.get_existing_output_files(with_suffix=file_type)
-
-    if not overwrite and (len(existing_output_files) > 0):
-        log.warning(
-            (
-                f"{existing_output_files} already exist(s)."
-                "To overwrite, set overwrite to True"
-            )
-        )
-        return True
-
-    return False
-
-
-def tracker_metadata(
-    sigma_l: float, sigma_h: float, sigma_iou: float, t_min: float, t_miss_max: float
-) -> dict:
-    return {
-        dataformat.NAME: "IOU",
-        dataformat.SIGMA_L: sigma_l,
-        dataformat.SIGMA_H: sigma_h,
-        dataformat.SIGMA_IOU: sigma_iou,
-        dataformat.T_MIN: t_min,
-        dataformat.T_MISS_MAX: t_miss_max,
-    }
-
-
-def mark_and_write_results(
-    chunk: TrackedChunk,
-    result_store: TrackingResultStore,
-    overwrite: bool,
-) -> None:
-    # no active tracks remaining, so last track frame metadata
-    # should be correct for all contained tracks,
-    # thus set finished flags now
-    mark_last_detections_as_finished(chunk, result_store)
-
-    # write marked detections to track file and delete the data
-    file_path = chunk.file_path
-
-    # reindex frames before writing ottrk file
-    file_type = CONFIG[DEFAULT_FILETYPE][TRACK]
-    serializable_detections: list[dict] = FrameIndexer().reindex(
-        chunk._detections, frame_offset=chunk.frame_offset
+    iou_tracker: IouTracker[Path] = IouTracker(
+        parameters=IouParameters(sigma_l, sigma_h, sigma_iou, t_min, t_miss_max)
     )
 
-    output = build_output(
-        serializable_detections,
-        chunk.metadata,
-        result_store.tracking_run_id,
-        chunk._frame_group_id,
-    )
-    write_json(
-        dict_to_write=output,
-        file=Path(file_path),
-        filetype=file_type,
-        overwrite=overwrite,
+    chunk_parser = JsonChunkParser()
+    group_parser = TimeThresholdFrameGroupParser(
+        tracker_data=tracker_metadata(sigma_l, sigma_h, sigma_iou, t_min, t_miss_max)
     )
 
-    log.info(f"Successfully tracked and wrote {file_path}")
-    del chunk._detections
-
-
-def track(
-    detections: dict,  # TODO: Type hint nested dict during refactoring
-    sigma_l: float = CONFIG[TRACK][IOU][SIGMA_L],
-    sigma_h: float = CONFIG[TRACK][IOU][SIGMA_H],
-    sigma_iou: float = CONFIG[TRACK][IOU][SIGMA_IOU],
-    t_min: int = CONFIG[TRACK][IOU][T_MIN],
-    t_miss_max: int = CONFIG[TRACK][IOU][T_MISS_MAX],
-    previous_active_tracks: list = [],
-    vehicle_id_generator: Iterator[int] = id_generator(),
-) -> TrackingResult:  # TODO: Type hint nested dict during refactoring
-    """Perform tracking using track_iou with arguments and add metadata to tracks.
-
-    Args:
-        detections (dict): Dict of detections in .otdet format.
-        sigma_l (float, optional): Lower confidence threshold. Detections with
-            confidences below sigma_l are not even considered for tracking.
-            Defaults to CONFIG["TRACK"]["IOU"]["SIGMA_L"].
-        sigma_h (float, optional): Upper confidence threshold. Tracks are only
-            considered as valid if they contain at least one detection with a confidence
-            above sigma_h.
-            Defaults to CONFIG["TRACK"]["IOU"]["SIGMA_H"].
-        sigma_iou (float, optional): Intersection-Over-Union threshold. Two detections
-            in subsequent frames are considered to belong to the same track if their IOU
-            value exceeds sigma_iou and this is the highest IOU of all possible
-            combination of detections.
-            Defaults to CONFIG["TRACK"]["IOU"]["SIGMA_IOU"].
-        t_min (int, optional): Minimum number of detections to count as a valid track.
-            All tracks with less detections will be dissmissed.
-            Defaults to CONFIG["TRACK"]["IOU"]["T_MIN"].
-        t_miss_max (int, optional): Maximum number of missed detections before
-            continuing a track. If more detections are missing, the track will not be
-            continued.
-            Defaults to CONFIG["TRACK"]["IOU"]["T_MISS_MAX"].
-        previous_active_tracks (list): a list of remaining active tracks
-            from previous iterations.
-        vehicle_id_generator (Iterator[int]): provides ids for new tracks
-
-    Returns:
-        TrackingResult: A result object holding tracked detecktions in ottrk format
-            and list of active tracks (iou format?)
-            and a lookup table for each tracks last detection frame.
-    """
-
-    result = track_iou(
-        detections=detections[DATA],
-        sigma_l=sigma_l,
-        sigma_h=sigma_h,
-        sigma_iou=sigma_iou,
-        t_min=t_min,
-        t_miss_max=t_miss_max,
-        previous_active_tracks=previous_active_tracks,
-        vehicle_id_generator=vehicle_id_generator,
+    file_tracker = GroupedFilesTracker(
+        tracker=iou_tracker,
+        chunk_parser=chunk_parser,
+        frame_group_parser=group_parser,
+        id_generator_factory=lambda _: track_id_generator(),
+        overwrite=True,
     )
-    log.info("Detections tracked")
-
-    return result
-
 
-def mark_last_detections_as_finished(
-    chunk: TrackedChunk, result_store: TrackingResultStore
-) -> None:
-    # invert last occurrence frame dict
-    last_track_frame = result_store.last_track_frame
-
-    frame_ending_tracks = defaultdict(set)
-    for vehID in chunk._detected_ids:
-        frame_ending_tracks[last_track_frame[vehID]].add(vehID)
+    buffer = UnfinishedChunksBuffer(
+        tracker=file_tracker,
+        keep_discarded=True,
+    )
 
-    for frame_num, frame_det in chunk._detections.items():
-        for ending_track in frame_ending_tracks[int(frame_num)]:
-            frame_det[ending_track][FINISHED] = True
-            del last_track_frame[ending_track]
+    exporter = FinishedChunkTrackExporter()
 
+    tracking_run_id = tracking_run_id_generator()
+    finished_chunk_stream = buffer.group_and_track(detections_files)
 
-def build_output(
-    detections: list[dict],
-    metadata: dict,
-    tracking_run_id: str,
-    frame_group_id: int,
-) -> dict:
-    metadata[dataformat.TRACKING][dataformat.TRACKING_RUN_ID] = tracking_run_id
-    metadata[dataformat.TRACKING][dataformat.FRAME_GROUP] = frame_group_id
-    return {METADATA: metadata, DATA: {DETECTIONS: detections}}
+    finished_chunk_progress = tqdm(
+        finished_chunk_stream, desc="export FrameChunk", total=len(detections_files)
+    )
+    exporter.export(tracking_run_id, iter(finished_chunk_progress), overwrite)

OTVision/track/tracker/filebased_tracking.py

@@ -0,0 +1,197 @@
+import logging
+from pathlib import Path
+from typing import Callable, Iterator
+
+from more_itertools import peekable
+from tqdm import tqdm
+
+from OTVision.config import CONFIG, DEFAULT_FILETYPE, OVERWRITE, TRACK
+from OTVision.helpers.log import LOGGER_NAME
+from OTVision.track.model.filebased.frame_chunk import (
+    ChunkParser,
+    FinishedChunk,
+    FrameChunk,
+    TrackedChunk,
+)
+from OTVision.track.model.filebased.frame_group import FrameGroup, FrameGroupParser
+from OTVision.track.model.frame import (
+    Frame,
+    FrameNo,
+    IsLastFrame,
+    TrackedFrame,
+    TrackId,
+)
+from OTVision.track.model.tracking_interfaces import (
+    ID_GENERATOR,
+    Tracker,
+    UnfinishedTracksBuffer,
+)
+
+log = logging.getLogger(LOGGER_NAME)
+
+
+class ChunkBasedTracker(Tracker[Path]):
+
+    def __init__(self, tracker: Tracker[Path], chunkParser: ChunkParser) -> None:
+        super().__init__()
+        self._chunk_parser = chunkParser
+        self._tracker = tracker
+
+    def track_frame(
+        self,
+        frames: Frame[Path],
+        id_generator: ID_GENERATOR,
+    ) -> TrackedFrame[Path]:
+        return self._tracker.track_frame(frames, id_generator)
+
+    def track_chunk(
+        self,
+        chunk: FrameChunk,
+        is_last_chunk: bool,
+        id_generator: ID_GENERATOR,
+    ) -> TrackedChunk:
+        frames_progress = tqdm(
+            chunk.frames, desc="track Frame", total=len(chunk.frames), leave=False
+        )
+
+        tracked_frames = self.track(iter(frames_progress), id_generator)
+        return TrackedChunk(
+            file=chunk.file,
+            frames=list(tracked_frames),
+            metadata=chunk.metadata,
+            is_last_chunk=is_last_chunk,
+            frame_group_id=chunk.frame_group_id,
+        )
+
+    def track_file(
+        self,
+        file: Path,
+        frame_group: FrameGroup,
+        is_last_file: bool,
+        id_generator: ID_GENERATOR,
+        frame_offset: int = 0,
+    ) -> TrackedChunk:
+        chunk = self._chunk_parser.parse(file, frame_group, frame_offset)
+        return self.track_chunk(chunk, is_last_file, id_generator)
+
+
+ID_GENERATOR_FACTORY = Callable[[FrameGroup], ID_GENERATOR]
+
+
+class GroupedFilesTracker(ChunkBasedTracker):
+
+    def __init__(
+        self,
+        tracker: Tracker[Path],
+        chunk_parser: ChunkParser,
+        frame_group_parser: FrameGroupParser,
+        id_generator_factory: ID_GENERATOR_FACTORY,
+        overwrite: bool = CONFIG[TRACK][OVERWRITE],
+        file_type: str = CONFIG[DEFAULT_FILETYPE][TRACK],
+    ) -> None:
+        super().__init__(tracker, chunk_parser)
+        self._group_parser = frame_group_parser
+        self._id_generator_of = id_generator_factory
+        self._overwrite = overwrite
+        self._file_type = file_type
+
+    def track_group(self, group: FrameGroup) -> Iterator[TrackedChunk]:
+        if self.check_skip_due_to_existing_output_files(group):
+            log.warning(f"Skip FrameGroup {group.id}")
+            yield from []  # TODO how to create empty generator stream?
+
+        frame_offset = 0  # frame no starts a 0 for each frame group
+        id_generator = self._id_generator_of(group)  # new id generator per group
+        file_stream = peekable(
+            tqdm(
+                group.files,
+                desc="track FrameChunk",
+                total=len(group.files),
+                leave=False,
+            )
+        )
+
+        for file in file_stream:
+            is_last = file_stream.peek(default=None) is None
+
+            chunk = self._chunk_parser.parse(file, group, frame_offset)
+            frame_offset = chunk.frames[-1].no + 1  # assuming frames are sorted by no
+
+            tracked_chunk = self.track_chunk(chunk, is_last, id_generator)
+            yield tracked_chunk
+
+    def group_and_track_files(self, files: list[Path]) -> Iterator[TrackedChunk]:
+        processed = self._group_parser.process_all(files)
+
+        processed_progress = tqdm(
+            processed, desc="track FrameGroup", total=len(processed), leave=False
+        )
+        for group in processed_progress:
+            yield from self.track_group(group)
+
+    def check_skip_due_to_existing_output_files(self, group: FrameGroup) -> bool:
+        if not self._overwrite and group.check_any_output_file_exists(self._file_type):
+            existing_files = group.get_existing_output_files(
+                with_suffix=self._file_type
+            )
+            log.warning(
+                (
+                    f"{existing_files} already exist(s)."
+                    "To overwrite, set overwrite to True"
+                )
+            )
+            return True
+
+        return False
+
+
+class UnfinishedChunksBuffer(UnfinishedTracksBuffer[TrackedChunk, FinishedChunk]):
+
+    def __init__(
+        self,
+        tracker: GroupedFilesTracker,
+        keep_discarded: bool = False,
+    ) -> None:
+        super().__init__(keep_discarded)
+        self.tracker = tracker
+
+    def group_and_track(self, files: list[Path]) -> Iterator[FinishedChunk]:
+        processed = self.tracker._group_parser.process_all(files)
+
+        processed_progress = tqdm(
+            processed, desc="track FrameGroup", total=len(processed), leave=False
+        )
+        for group in processed_progress:
+            yield from self.track_group(group)
+
+    def track_group(self, group: FrameGroup) -> Iterator[FinishedChunk]:
+        tracked_chunk_stream = self.tracker.track_group(group)
+        return self.track_and_finish(tracked_chunk_stream)
+
+    def _get_last_track_frames(self, container: TrackedChunk) -> dict[TrackId, FrameNo]:
+        return container.last_track_frame
+
+    def _get_unfinished_tracks(self, container: TrackedChunk) -> set[TrackId]:
+        return container.unfinished_tracks
+
+    def _get_observed_tracks(self, container: TrackedChunk) -> set[TrackId]:
+        return container.observed_tracks
+
+    def _get_newly_finished_tracks(self, container: TrackedChunk) -> set[TrackId]:
+        return container.finished_tracks
+
+    def _get_newly_discarded_tracks(self, container: TrackedChunk) -> set[TrackId]:
+        return container.discarded_tracks
+
+    def _get_last_frame_of_container(self, container: TrackedChunk) -> FrameNo:
+        return max(frame.no for frame in container.frames)
+        # todo faster implementation if sorted or save as metadata?
+
+    def _finish(
+        self,
+        container: TrackedChunk,
+        is_last: IsLastFrame,
+        discarded_tracks: set[TrackId],
+        keep_discarded: bool,
+    ) -> FinishedChunk:
+        return container.finish(is_last, discarded_tracks, keep_discarded)