OTVision 0.6.1__py3-none-any.whl → 0.6.3__py3-none-any.whl

OTVision/track/model/filebased/frame_chunk.py

@@ -0,0 +1,203 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Sequence
+
+from tqdm import tqdm
+
+from OTVision.dataformat import FRAME, INPUT_FILE_PATH, TRACK_ID
+from OTVision.domain.detection import TrackId
+from OTVision.domain.frame import (
+    DetectedFrame,
+    FinishedFrame,
+    FrameNo,
+    IsLastFrame,
+    TrackedFrame,
+)
+from OTVision.track.model.filebased.frame_group import FrameGroup, get_output_file
+
+
+@dataclass(frozen=True)
+class FrameChunk:
+    """
+    A chunk of Frames obtained from a common file path source.
+
+    Attributes:
+        file (Path): common file path source of Frames.
+        metadata (dict): otdet metadata.
+        frames (Sequence[DetectedFrame]): a sequence of untracked Frames.
+        frame_group_id (int): id of FrameGroup this FrameCHunk is part of.
+    """
+
+    file: Path
+    metadata: dict
+    frames: Sequence[DetectedFrame]
+    frame_group_id: int
+
+    def check_output_file_exists(self, with_suffix: str) -> bool:
+        return get_output_file(self.file, with_suffix).is_file()
+
+
+@dataclass(frozen=True)
+class TrackedChunk(FrameChunk):
+    """A chunk of TrackedFrames extends FrameChunk.
+    Aggregates observed/finished/unfinished tracking metadata.
+
+    If is_last_chunk is true, all unfinished tracks
+    are marked as finished in last frame.
+
+    Attributes:
+        is_last_chunk (bool): whether this chunk is the last of
+            subsequently (related/connected) chunks.
+        frames (Sequence[TrackedFrame]): overrides frames
+            with more specific frame type.
+        finished_tracks (set[TrackId]): aggregates finished tracks
+            of given TrackedFrames.
+        observed_tracks (set[TrackId]): aggregates observed tracks
+            of given TrackedFrames.
+        unfinished_tracks (set[TrackId]): aggregates unfinished tracks
+            of given TrackedFrames as observed but not finished tracks in chunk.
+        last_track_frame (dict[TrackId, int]): mapping of track id
+            to frame number in which it last occurs.
+    """
+
+    is_last_chunk: bool
+    frames: Sequence[TrackedFrame] = field(init=False)
+
+    finished_tracks: set[TrackId] = field(init=False)
+    observed_tracks: set[TrackId] = field(init=False)
+    unfinished_tracks: set[TrackId] = field(init=False)
+    discarded_tracks: set[TrackId] = field(init=False)
+    last_track_frame: dict[TrackId, int] = field(init=False)
+
+    def __init__(
+        self,
+        file: Path,
+        metadata: dict,
+        is_last_chunk: bool,
+        frames: Sequence[TrackedFrame],
+        frame_group_id: int,
+    ) -> None:
+
+        object.__setattr__(self, "file", file)
+        object.__setattr__(self, "metadata", metadata)
+        object.__setattr__(self, "is_last_chunk", is_last_chunk)
+        object.__setattr__(self, "frame_group_id", frame_group_id)
+
+        observed = set().union(*(f.observed_tracks for f in frames))
+        finished = set().union(
+            *(f.finished_tracks for f in frames)
+        )  # TODO remove discarded?
+        discarded = set().union(*(f.discarded_tracks for f in frames))
+        unfinished = {o for o in observed if o not in finished and o not in discarded}
+
+        # set all unfinished tracks as finished, as this is the last track
+        if self.is_last_chunk:
+            frames_list = list(frames)
+            # assume frames sorted by occurrence
+
+            last_frame = frames_list[-1]
+            frames_list[-1] = replace(
+                last_frame,
+                finished_tracks=last_frame.finished_tracks.union(unfinished),
+            )
+
+            unfinished = set()
+            finished = set().union(*(f.finished_tracks for f in frames_list))
+        else:
+            frames_list = list(frames)
+
+        object.__setattr__(self, "frames", frames_list)
+
+        object.__setattr__(self, "finished_tracks", finished)
+        object.__setattr__(self, "observed_tracks", observed)
+        object.__setattr__(self, "unfinished_tracks", unfinished)
+        object.__setattr__(self, "discarded_tracks", discarded)
+
+        # assume frames sorted by occurrence
+        last_track_frame: dict[TrackId, FrameNo] = {
+            detection.track_id: frame.no
+            for frame in self.frames
+            for detection in frame.detections
+        }
+        object.__setattr__(self, "last_track_frame", last_track_frame)
+
+    def finish(
+        self,
+        is_last: IsLastFrame,
+        discarded_tracks: set[TrackId],
+        keep_discarded: bool = False,
+    ) -> "FinishedChunk":
+        """Turn this TrackedChunk into a FinishedChunk
+        by adding is_finished information to all its detections.
+
+        Args:
+            is_last (IsLastFrame): function to determine whether
+                a track is finished in a certain frame.
+            discarded_tracks (set[TrackId]): list of tracks considered discarded.
+                Used to mark corresponding tracks.
+            keep_discarded (bool): whether FinishedDetections marked as discarded
+                should be kept in detections list. Defaults to False.
+
+        Returns:
+            FinishedChunk: chunk of FinishedFrames
+        """
+        return FinishedChunk(
+            file=self.file,
+            metadata=self.metadata,
+            is_last_chunk=self.is_last_chunk,
+            frames=[
+                frame.finish(is_last, discarded_tracks, keep_discarded)
+                for frame in self.frames
+            ],
+            frame_group_id=self.frame_group_id,
+        )
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def __str__(self) -> str:
+        return f"FG [{self.frame_group_id}] - {self.file}"
+
+
+@dataclass(frozen=True)
+class FinishedChunk(TrackedChunk):
+    """A chunk of FinishedFrames.
+
+    Attributes:
+        frames (Sequence[FinishedFrame]): overrides frames
+            with more specific frame type.
+    """
+
+    frames: Sequence[FinishedFrame]
+
+    def to_detection_dicts(self) -> list[dict]:
+        chunk_metadata = {INPUT_FILE_PATH: self.file.as_posix()}
+
+        frames_progress = tqdm(
+            self.frames, desc="Frames to_dict", total=len(self.frames), leave=False
+        )
+
+        detection_dict_list = [
+            {**det_dict, **chunk_metadata}
+            for frame in frames_progress
+            for det_dict in frame.to_detection_dicts()
+        ]
+
+        detection_dict_list.sort(
+            key=lambda detection: (
+                detection[FRAME],
+                detection[TRACK_ID],
+            )
+        )
+        return detection_dict_list
+
+
+class ChunkParser(ABC):
+    """A parser for file path to FrameChunk."""
+
+    @abstractmethod
+    def parse(
+        self, file: Path, frame_group: FrameGroup, frame_offset: int
+    ) -> FrameChunk:
+        pass

OTVision/track/model/filebased/frame_group.py

@@ -0,0 +1,95 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, replace
+from datetime import datetime
+from pathlib import Path
+
+from tqdm import tqdm
+
+
+def get_output_file(file: Path, with_suffix: str) -> Path:
+    return file.with_suffix(with_suffix)
+
+
+@dataclass(frozen=True)
+class FrameGroup:
+    id: int
+    start_date: datetime
+    end_date: datetime
+    hostname: str
+    files: list[Path]
+    metadata_by_file: dict[Path, dict]
+
+    def merge(self, other: "FrameGroup") -> "FrameGroup":
+        if self.start_date < other.start_date:
+            return FrameGroup._merge(self, other)
+        else:
+            return FrameGroup._merge(other, self)
+
+    @staticmethod
+    def _merge(first: "FrameGroup", second: "FrameGroup") -> "FrameGroup":
+        if first.hostname != second.hostname:
+            raise ValueError("Hostname of FrameGroups does not match")
+
+        files = first.files + second.files
+        metadata = dict(first.metadata_by_file)
+        metadata.update(second.metadata_by_file)
+
+        merged = FrameGroup(
+            id=first.id,
+            start_date=first.start_date,
+            end_date=second.end_date,
+            hostname=first.hostname,
+            files=files,
+            metadata_by_file=metadata,
+        )
+
+        return merged
+
+    def check_any_output_file_exists(self, with_suffix: str) -> bool:
+        return len(self.get_existing_output_files(with_suffix)) > 0
+
+    def get_existing_output_files(self, with_suffix: str) -> list[Path]:
+        return [file for file in self.get_output_files(with_suffix) if file.is_file()]
+
+    def get_output_files(self, with_suffix: str) -> list[Path]:
+        return [get_output_file(file, with_suffix) for file in self.files]
+
+    def with_id(self, new_id: int) -> "FrameGroup":
+        return replace(self, id=new_id)
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def __str__(self) -> str:
+        return f"FrameGroup[{self.id}] = [{self.start_date} - {self.end_date}]"
+
+
+class FrameGroupParser(ABC):
+
+    def process_all(self, files: list[Path]) -> list[FrameGroup]:
+        files_progress = tqdm(files, desc="parse FrameGroups", total=len(files))
+
+        parsed: list[FrameGroup] = [self.parse(file) for file in files_progress]
+        merged: list[FrameGroup] = self.merge(parsed)
+        updated: list[FrameGroup] = [
+            self.updated_metadata_copy(group).with_id(i)
+            for i, group in enumerate(merged)
+        ]
+
+        return updated
+
+    @abstractmethod
+    def parse(self, file: Path) -> FrameGroup:
+        pass
+
+    @abstractmethod
+    def merge(self, frame_groups: list[FrameGroup]) -> list[FrameGroup]:
+        pass
+
+    def updated_metadata_copy(self, frame_group: FrameGroup) -> FrameGroup:
+        new_metadata = self.update_metadata(frame_group)
+        return replace(frame_group, metadata_by_file=new_metadata)
+
+    @abstractmethod
+    def update_metadata(self, frame_group: FrameGroup) -> dict[Path, dict]:
+        pass

OTVision/track/model/track_exporter.py

@@ -0,0 +1,119 @@
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Generic, Iterator, TypeVar
+
+from tqdm import tqdm
+
+from OTVision.config import CONFIG, DEFAULT_FILETYPE, TRACK
+from OTVision.dataformat import (
+    DATA,
+    DETECTIONS,
+    FRAME,
+    FRAME_GROUP,
+    INPUT_FILE_PATH,
+    METADATA,
+    TRACK_ID,
+    TRACKING,
+    TRACKING_RUN_ID,
+)
+from OTVision.helpers.files import write_json
+from OTVision.helpers.log import LOGGER_NAME
+
+log = logging.getLogger(LOGGER_NAME)
+
+
+F = TypeVar("F")  # Finished container: e.g. FinishedFrame or FinishedChunk
+
+
+class FinishedTracksExporter(ABC, Generic[F]):
+
+    def __init__(self, file_type: str = CONFIG[DEFAULT_FILETYPE][TRACK]):
+        self.file_type = file_type
+
+    @abstractmethod
+    def get_detection_dicts(self, container: F) -> list[dict]:
+        pass
+
+    @abstractmethod
+    def get_result_path(self, container: F) -> Path:
+        pass
+
+    @abstractmethod
+    def get_metadata(self, container: F) -> dict:
+        pass
+
+    @abstractmethod
+    def get_frame_group_id(self, container: F) -> int:
+        pass
+
+    def export(
+        self, tracking_run_id: str, stream: Iterator[F], overwrite: bool
+    ) -> None:
+        for container in stream:
+            self.export_frames(container, tracking_run_id, overwrite)
+
+    def export_frames(
+        self, container: F, tracking_run_id: str, overwrite: bool
+    ) -> None:
+        file_path = self.get_result_path(container)
+
+        det_dicts = self.reindex(self.get_detection_dicts(container))
+
+        output = self.build_output(
+            det_dicts,
+            self.get_metadata(container),
+            tracking_run_id,
+            self.get_frame_group_id(container),
+        )
+
+        write_json(
+            dict_to_write=output,
+            file=Path(file_path),
+            filetype=self.file_type,
+            overwrite=overwrite,
+        )
+
+        log.info(f"Successfully tracked and wrote {file_path}")
+
+    @staticmethod
+    def reindex(det_dicts: list[dict]) -> list[dict]:
+        min_frame_no = min(det[FRAME] for det in det_dicts)
+
+        det_dicts_progress = tqdm(
+            det_dicts,
+            desc="reindex TrackedDetections",
+            total=len(det_dicts),
+            leave=False,
+        )
+        reindexed_dets = [
+            {**det, **{FRAME: det[FRAME] - min_frame_no + 1}}
+            for det in det_dicts_progress
+        ]
+
+        if len(reindexed_dets) == 0:
+            return []
+
+        if len({detection[INPUT_FILE_PATH] for detection in reindexed_dets}) > 1:
+            raise ValueError("Expect detections from only a single source file")
+
+        reindexed_dets.sort(
+            key=lambda detection: (
+                detection[INPUT_FILE_PATH],
+                detection[FRAME],
+                detection[TRACK_ID],
+            )
+        )
+
+        return reindexed_dets
+
+    @staticmethod
+    def build_output(
+        detections: list[dict],
+        metadata: dict,
+        tracking_run_id: str,
+        frame_group_id: int,
+    ) -> dict:
+        metadata[TRACKING][TRACKING_RUN_ID] = tracking_run_id
+        metadata[TRACKING][FRAME_GROUP] = frame_group_id
+        return {METADATA: metadata, DATA: {DETECTIONS: detections}}