Data collection by temporal consistency of object detection and tracking
# config.py

# Directories or paths related to tests
TEST_DIRS = ["tests_unit"]

# Coverage settings
COV_PATHS = ["temporal_consistency"]
COV_OPTIONS = {
    "branch": True,
    "report": "term-missing",
    "fail_under": 5,
}

# Add any custom fixtures or settings as needed for your pytest setup

config-3.py
smttsp/temporal_consistency_odt
# This should be left blank
# This is a support file for `pytest` to allow importing of the main package

conftest.py
smttsp/temporal_consistency_odt
"""This script is an entrypoint for performing object detection and tracking
on a given video using YOLO and Deep Sort. It also includes options for
augmenting frames and adjusts tracking based on configurable settings like
confidence threshold and max age for tracks.

Detected temporal anomalies in tracking are logged for further analysis.
They will later be used for model training and fine-tuning.
"""

import os

import configargparse
from deep_sort_realtime.deepsort_tracker import DeepSort
from loguru import logger
from ultralytics import YOLO

from temporal_consistency.frame_anomaly_detection import TemporalAnomalyDetector
from temporal_consistency.object_detection_tracking import (
    run_detection_and_tracking_pipeline,
)
from temporal_consistency.utils import get_runtime_str


CONFIDENCE_THRESHOLD = 0.4
MAX_AGE = 25


def parse_args():
    parser = configargparse.ArgumentParser(
        description="Demo for parsing different types of data"
    )
    parser.add_argument(
        "--config",
        is_config_file=True,
        default="conf.yaml",
        help="Path to the config file",
    )
    parser.add_argument(
        "--video_filepath",
        help="Path to the input video file",
    )
    parser.add_argument(
        "--out_folder",
        help="Path to the output folder",
    )
    parser.add_argument(
        "--confidence",
        type=float,
        default=CONFIDENCE_THRESHOLD,
        help="Filtering predictions with low confidence",
    )
    parser.add_argument(
        "--max_age",
        type=float,
        default=MAX_AGE,
        help="Filtering predictions with low confidence",
    )
    parser.add_argument(
        "--details",
        type=int,
        default=1,
        choices=[1, 2, 3],
        help="Indicates how much details you want in analysis, "
        "1: minimal details"
        "2: more details"
        "3: most details",
    )
    parser.add_argument(
        "--num_aug",
        type=int,
        default=0,
        choices=[0, 1, 2, 3],
        help="Each frame will go through up to 3 augmentations. 0-> no augmentation",
    )
    parser.add_argument(
        "--out_video_fps",
        type=int,
        default=25,
        help="FPS of the output videos",
    )

    args = parser.parse_args()
    return args


def main(args):
    runtime_str = get_runtime_str()
    args.out_folder = os.path.join(args.out_folder, runtime_str)
    os.makedirs(args.out_folder, exist_ok=True)

    logfile = os.path.join(args.out_folder, "output.log")
    logger.add(logfile)

    model = YOLO("yolov8n.pt")
    deep_sort_tracker = DeepSort(max_age=args.max_age)

    tframe_collection = run_detection_and_tracking_pipeline(
        model, deep_sort_tracker, args
    )
    TemporalAnomalyDetector(tframe_collection)


if __name__ == "__main__":
    args = parse_args()

    main(args)

main.py
smttsp/temporal_consistency_odt
"""The `temporal_consistency` Python Package

This is an example Python package for use by Samet looking
to build out new Python packages
"""

temporal_consistency/__init__.py
smttsp/temporal_consistency_odt
"""This module provides functions to generate a list of image augmentations
using the albumentations library and to apply a random subset of those
augmentations to a given image. Depending on the `num_aug` input, the
`get_random_augmentation` function selects a specified number of augmentations
from this list of possible options and applies them to an input image.
"""

import random

import numpy
from albumentations import (
    Blur,
    ChannelShuffle,
    ColorJitter,
    Compose,
    Equalize,
    GaussNoise,
    InvertImg,
    Posterize,
    RandomBrightnessContrast,
    RandomFog,
    RandomGamma,
    RandomRain,
    RandomShadow,
    RandomSnow,
    RandomSunFlare,
    RandomToneCurve,
    RGBShift,
    Solarize,
)


def get_aug_list():
    """Returns a list of augmentations to be applied to the frames."""

    return [
        RandomBrightnessContrast(p=1.0),
        RandomGamma(p=1.0),
        ColorJitter(
            brightness=random.uniform(0.1, 0.4),
            contrast=random.uniform(0.1, 0.5),
            saturation=random.uniform(0.1, 0.5),
            hue=random.uniform(0.1, 0.5),
            p=1.0,
        ),
        ChannelShuffle(p=1.0),
        RGBShift(p=1.0),
        Blur(blur_limit=3, p=1.0),
        GaussNoise(p=1.0),
        RandomRain(p=1.0),
        RandomSnow(p=1.0),
        RandomSunFlare(p=1.0, flare_roi=(0, 0, 1, 0.5)),
        RandomFog(p=1.0),
        RandomShadow(p=1.0),
        InvertImg(p=1.0),
        RandomToneCurve(p=1.0),
        Solarize(p=1.0),
        Equalize(p=1.0),
        Posterize(p=1.0),
    ]


def get_random_augmentation(image: numpy.ndarray, num_aug=0):
    """Pick random augmentations from the list and apply them to the image.
    The number of augmentations to apply is specified by num_aug.
    """

    image_aug = image
    if num_aug > 0:
        augmentation_pipeline = Compose(
            random.sample(get_aug_list(), k=num_aug)
        )
        image_aug = augmentation_pipeline(image=image)["image"]
    return image_aug

temporal_consistency/augmentations.py
smttsp/temporal_consistency_odt
"""This module contains `TemporalAnomalyDetector` class for identifying inconsistencies
in object tracking across a sequence of frames. It checks for issues like classification
inconsistencies, missing objects in frames, single-frame appearances, and low
Intersection-over-Union (IoU) values. Detected anomalies are stored in a dictionary.
"""
import os
import sys
from collections import defaultdict

import cv2
from loguru import logger

from temporal_consistency.tracked_frame import TrackedFrameCollection
from temporal_consistency.utils import compute_iou
from temporal_consistency.vis_utils import draw_class_name


MIN_IOU_THRESH = 0.5
EPS = sys.float_info.epsilon


def export_list_of_objects(object_filepath, object_list):
    with open(object_filepath, "a") as f:
        for obj in object_list:
            f.write(f"{obj.to_str()}\n")
        f.write("\n")


class TemporalAnomalyDetector:
    """Detects anomalies in the temporal consistency of the tracked objects."""

    def __init__(self, tframe_collection: TrackedFrameCollection):
        """Initializes the TemporalAnomalyDetector.

        Args:
            tframe_collection (TrackedFrameCollection): A collection of frames
                containing tracked objects.
        """

        self.tframe_collection = tframe_collection
        self.anomalies: defaultdict = defaultdict(list)
        self.scan_for_anomalies()
        self.export_anomalies()

    def scan_for_anomalies(self):
        """Scans for anomalies across all objects in the frame collection."""

        for object_id, track_info in self.tframe_collection.all_objects.items():
            self.inspect_object_for_anomalies(object_id, track_info)

        return None

    def inspect_object_for_anomalies(
        self, object_id: str, track_info: dict
    ) -> bool:
        """Checks for potential anomalies for a single tracked object.

        Args:
            object_id (str): Unique identifier of the tracked object.
            track_info (dict): A dictionary containing frame IDs as keys and prediction
                objects as values, where each prediction has a 'class_name' attribute.

        Returns:
            bool: True if anomalies are detected, False otherwise.
        """

        anomaly_exist = (
            self.is_class_inconsistent(object_id, track_info)
            or self.is_object_missing_in_frames(object_id, track_info)
            or self.appears_only_in_single_frame(object_id, track_info)
            or self.has_low_iou(object_id, track_info)
        )
        return anomaly_exist

    def is_class_inconsistent(self, object_id: str, track_info: dict) -> bool:
        """Verifies that an object maintains consistent classification across frames.

        Returns:
            bool: False if the class is consistent, True if there are inconsistencies.
        """

        all_classes = set(v.class_name for v in track_info.values())

        if len(all_classes) > 1:
            log = f"{object_id=} occurs as the following classes: {all_classes}"
            logger.info(log)
            self.get_frame_id_for_class_inconsistency(object_id, track_info)

        return len(all_classes) > 1

    def get_frame_id_for_class_inconsistency(self, object_id, track_info):
        """Find the first frame ID where a class inconsistency occurs in
            a tracking sequence.

        Returns:
            None: The function modifies the 'anomalies' attribute of the object,
                appending frame IDs where class inconsistency is detected.
        """

        prev_class = None

        for frame_id, pred in track_info.items():
            current_class = pred.class_name
            if prev_class is not None and prev_class != current_class:
                self.anomalies[object_id].append(frame_id)
                break

            prev_class = current_class

        return None

    def is_object_missing_in_frames(
        self, object_id: str, track_info: dict
    ) -> bool:
        """Checks if an object goes missing in intermediate frames.

        Returns:
            bool: True if the object is missing in some frames, False otherwise.
        """

        keys = track_info.keys()
        mn_idx, mx_idx, size = min(keys), max(keys), len(track_info)
        expected_size = mx_idx - mn_idx + 1
        if size != expected_size:
            log = f"{object_id=} is missing in {expected_size - size} frames"
            logger.info(log)

            # sampling only one frame where the object is missing
            frame_ids = set(range(mn_idx, mx_idx + 1)).difference(keys)
            self.anomalies[object_id].append(list(frame_ids)[0])

        return expected_size != size

    def appears_only_in_single_frame(
        self, object_id: str, track_info: dict
    ) -> bool:
        """Checks if an object only appears in a single frame, potentially
            indicating a false detection.

        Returns:
            bool: True if the object appears in only one frame, False otherwise.
        """

        if len(track_info) == 1:
            log = f"{object_id=} occurs only in one frame, may indicate false detection"
            logger.info(log)
            frame_id = list(track_info.keys())[0]
            self.anomalies[object_id].append(frame_id)

        return len(track_info) == 1

    def has_low_iou(self, object_id: str, track_info: dict) -> bool:
        """Assesses if the Intersection over Union (IoU) is below a threshold,
            indicating potential tracking issues.

        Returns:
            bool: True if there is low IoU, False otherwise.
        """

        keys = sorted(list(track_info.keys()))
        flag = False
        for frame_i, frame_j in zip(keys, keys[1:]):
            t1 = track_info[frame_i]
            t2 = track_info[frame_j]

            # TODO (samet): Need to check if the bboxes are close to frame edges
            iou = compute_iou(t1.ltrb, t2.ltrb)

            if iou < MIN_IOU_THRESH:
                log = (
                    f"{iou=} is lower than threshold of {MIN_IOU_THRESH} "
                    f"between {frame_i=} and {frame_j=}"
                )
                logger.info(log)
                self.anomalies[object_id].append(frame_i)
                flag = True

        return flag

    def export_anomalies(self):
        """Exports the frames where there is an anomaly for at least one object.
        The exported files are:
        1. Raw frame as frame{frame_id}.jpg
        2. Frame with bboxes as frame{frame_id}_bbox.jpg. The green bboxes are the
            tracked objects, and the red bboxes are the low-confidence detections.
        3. a file containing all the bboxes in the frame as frame{frame_id}_bbox.txt

        1 + 3 can be used to help with labeling the data (i.e. model assisted labeling)
        """

        frame_ids = set(x[0] for x in self.anomalies.values())

        out_folder = os.path.join(self.tframe_collection.out_folder)
        os.makedirs(out_folder, exist_ok=True)

        for frame_id in frame_ids:
            frame_filepath = os.path.join(out_folder, f"frame{frame_id}.jpg")
            object_filepath = frame_filepath.replace(".jpg", "_bbox.txt")
            frame_wbbox_filepath = frame_filepath.replace(".jpg", "_bbox.jpg")

            # 1. Raw frame
            frame = self.tframe_collection.get_frame(frame_id)

            cv2.imwrite(frame_filepath, frame)

            # 2. Frame with bboxes
            (
                high_conf_objects,
                low_conf_objects,
            ) = self.tframe_collection.get_frame_predictions(frame_id)
            export_list_of_objects(object_filepath, high_conf_objects)
            export_list_of_objects(object_filepath, low_conf_objects)

            # 3. Bboxes as txt
            for obj in high_conf_objects:
                draw_class_name(frame, obj.ltrb, "", obj.class_name)

            for obj in low_conf_objects:
                draw_class_name(
                    frame, obj.ltrb, "", obj.class_name, bbox_color=(0, 0, 255)
                )
            cv2.imwrite(frame_wbbox_filepath, frame)

        return None

temporal_consistency/frame_anomaly_detection.py
smttsp/temporal_consistency_odt
"""This module handles object detection and tracking in video sequences.
It performs

- real-time object detection using a YOLO model
- object tracking using Deep SORT.

The frames can be optionally augmented before processing which is
for helping with robustness of the object detection model (i.e., finding failures).

The output includes visualization of object tracking (one video with bboxes and
one video per object with the object's track).
"""

import datetime

import cv2
import numpy
import torch
from deep_sort_realtime.deepsort_tracker import DeepSort

from temporal_consistency.augmentations import get_random_augmentation
from temporal_consistency.tracked_frame import (
    TrackedFrame,
    TrackedFrameCollection,
)
from temporal_consistency.utils import create_video_writer
from temporal_consistency.vis_utils import (
    draw_bbox_around_object,
    draw_fps_on_frame,
)


def transform_detection_predictions(data: list) -> list:
    """Transforms raw detection data into a structured format.

    Args:
        data (list): List of raw data for a detected object. The expected format
            includes bounding box coordinates, confidence score, and class ID.

    Returns:
        list: Transformed data for the detected object in the format
            [ltwh, confidence, class_id].
    """

    x_min, y_min, x_max, y_max = map(int, data[:4])
    ltwh = [x_min, y_min, x_max - x_min, y_max - y_min]
    confidence = data[4]
    class_id = int(data[5])

    res = [ltwh, confidence, class_id]
    return res


def object_detection(
    model, frame: numpy.ndarray, num_aug=0, confidence_threshold=0.1
):
    """Performs object detection on the given frame and returns the results.

    Args:
        model (YOLO): Model used for object detection.
        frame (numpy.ndarray): Frame on which objects are detected.
        num_aug (int, optional): Number of augmentations to apply to the frame.
        confidence_threshold (float, optional): Threshold for object detection.
    """

    frame_aug = get_random_augmentation(frame, num_aug=num_aug)

    with torch.no_grad():
        detections = model(frame_aug)[0]

    all_results = [
        transform_detection_predictions(data)
        for data in detections.boxes.data.tolist()
    ]
    results = []
    low_confidence_results = []

    for res in all_results:
        if res[1] >= confidence_threshold:
            results.append(res)
        else:
            low_confidence_results.append(res)

    return results, low_confidence_results, frame_aug


def object_tracking(
    frame: numpy.ndarray,
    results: list,
    deep_sort_tracker: DeepSort,
    classes: dict,
) -> numpy.ndarray:
    """Processes the given frame with object tracking using Deep SORT
        and visualizes the tracking results.

    The function takes an input frame, object detection results, a DeepSort
    tracker instance, and a dictionary mapping class IDs to class names.
    It updates the tracker with the new detection results and draws bounding boxes
    around the confirmed tracks on a copy of the input frame.

    Args:
        frame (numpy.ndarray): The frame on which objects are detected and tracked.
        results (list): List of object detection results for the given frame.
        deep_sort_tracker (DeepSort): Instance of the DST to update and track objects.
        classes (dict): Dictionary mapping class IDs to class names for visualization.

    Returns:
        numpy.ndarray: Frame with drawn bboxes around the confirmed tracked objects.
    """

    tracks = deep_sort_tracker.update_tracks(results, frame=frame)

    frame_after = frame.copy()
    for track in tracks:
        if not track.is_confirmed():
            continue

        voc_bbox = track.to_ltrb()
        draw_bbox_around_object(frame_after, track, voc_bbox, classes)
    return frame_after


def process_single_frame(
    model,
    video_cap: cv2.VideoCapture,
    frame_id: int,
    tframe_collection: TrackedFrameCollection,
    deep_sort_tracker: DeepSort,
    num_aug: int,
    confidence_threshold: float,
):
    """Processes a single frame from the video. This function does object
        detection and tracking. It also updates the TrackedFrameCollection.
        The last step is to draw the FPS on the frame.

    Args:
        model (YOLO): Model used for object detection.
        video_cap (cv2.VideoCapture): Video capture object.
        frame_id (int): Frame ID.
        tframe_collection (TrackedFrameCollection): Collection of tracked frames.
        deep_sort_tracker (DeepSort): Deep SORT tracker.
        num_aug (int): Number of augmentations to apply to the frame.
        confidence_threshold (float): Confidence threshold for object detection.

    Returns:
        numpy.ndarray: Frame with drawn bboxes around the confirmed tracked objects.
    """

    start = datetime.datetime.now()

    ret, frame = video_cap.read()
    if not ret:
        return None

    results, low_confidence_results, frame_aug = object_detection(
        model, frame, num_aug, confidence_threshold
    )
    frame_after = object_tracking(
        frame_aug, results, deep_sort_tracker, classes=model.names
    )
    tframe = TrackedFrame(
        frame_id,
        frame_aug,
        deep_sort_tracker.tracker,
        low_confidence_results,
        class_names=model.names,
    )
    tframe_collection.add_tracked_frame(tframe)
    end = datetime.datetime.now()

    total_time = (end - start).total_seconds() * 1000
    draw_fps_on_frame(frame_after, total_time)

    return frame_after


def apply_detection_and_tracking(
    model,
    deep_sort_tracker: DeepSort,
    num_aug: int,
    video_cap: cv2.VideoCapture,
    writer: cv2.VideoWriter,
    out_folder: str,
    out_video_fps: int,
    confidence_threshold: float,
) -> TrackedFrameCollection:
    """Applies object detection and tracking on video frames using
    the provided model and tracker.

    This function processes a video by detecting objects in its frames and
    then tracking those objects using Deep SORT. The results, including bboxes,
    are written to a video. The tracked objects are also saved as separate videos.

    Args:
        model (YOLO): Model used for object detection.
        deep_sort_tracker (DeepSort): Deep SORT tracker instance for object tracking.
        num_aug (int): Number of augmentations to apply to the frame.
        video_cap (cv2.VideoCapture): Video capture object to read frames from.
        writer (cv2.VideoWriter): Video writer object to output the processed video.
        out_folder (str): Output folder path where tracked objects will be saved.
        out_video_fps (int): Frames per second for the output video.
        confidence_threshold (float): Confidence threshold for object detection.

    Returns:
        TrackedFrameCollection: A collection of frames with tracking information.
    """

    tframe_collection = TrackedFrameCollection(
        video_cap=video_cap,
        class_names=model.names,
        out_folder=out_folder,
    )

    frame_id = 0
    while True:
        frame_after = process_single_frame(
            model,
            video_cap,
            frame_id,
            tframe_collection,
            deep_sort_tracker,
            num_aug,
            confidence_threshold,
        )
        if frame_after is None:
            break

        if frame_id >= 100:
            break

        writer.write(frame_after)
        frame_id += 1

    tframe_collection.export_all_objects(out_video_fps=out_video_fps)

    return tframe_collection


def run_detection_and_tracking_pipeline(
    model, deep_sort_tracker: DeepSort, args
):
    """Performs object detection and tracking on the given video.
    It also outputs the tracked objects into separate videos.

    Args:
        model (YOLO): Model used for object detection.
        deep_sort_tracker (DeepSort): Deep SORT tracker.
        args (argparse.Namespace): Command line arguments obtained from config file.

    Returns:
        TrackedFrameCollection: Collection of tracked frames.
    """

    video_filepath = args.video_filepath
    num_aug = args.num_aug
    confidence_threshold = args.confidence
    out_folder = args.out_folder
    out_video_fps = args.out_video_fps

    video_cap = cv2.VideoCapture(video_filepath)
    output_filepath = video_filepath.replace(".mp4", "_output.mp4")
    writer = create_video_writer(video_cap, output_filepath, fps=out_video_fps)

    tframe_collection = apply_detection_and_tracking(
        model,
        deep_sort_tracker,
        num_aug,
        video_cap,
        writer,
        out_folder,
        out_video_fps,
        confidence_threshold,
    )

    video_cap.release()
    writer.release()
    cv2.destroyAllWindows()

    return tframe_collection

temporal_consistency/object_detection_tracking.py
smttsp/temporal_consistency_odt
"""This module defines three classes: Prediction, TrackedFrame, TrackedFrameCollection
for managing object tracking in video frames.

- `Prediction` class encapsulates information about a detected object in a single frame,
including its bounding box, confidence score, and class.
- `TrackedFrame` class holds information for a single video frame, capturing all
its tracked objects along with some low-confidence detections.
- `TrackedFrameCollection` serves as a collection of TrackedFrames. It facilitates
operations such as adding new tracked frames to the collection, and exporting
the objects to individual videos.

Together, they provide a comprehensive structure for managing and exporting
object tracking data.
"""

import copy
import os
from collections import defaultdict

import cv2
import numpy

from temporal_consistency.utils import create_video_writer, ltwh_to_ltrb
from temporal_consistency.vis_utils import put_text_on_upper_corner


class Prediction:
    """Single prediction for an object. Contains the bounding box coordinates,
    confidence, class ID, and class name.
    """

    def __init__(
        self,
        frame_id: int,
        ltrb: list[int],
        confidence: float,
        class_id: int,
        class_names: dict,
    ):
        self.frame_id = frame_id
        self.ltrb = ltrb
        self.confidence = confidence
        self.class_id = class_id
        self.class_name = class_names.get(class_id, None)

    def to_str_all(self):
        confidence = round(self.confidence, 4) if self.confidence else None
        bbox = " ".join(map(str, self.ltrb))
        return (
            f"{self.frame_id=}, {self.class_id=}, {self.class_name=}, "
            f"{bbox}, {confidence=}"
        )

    def to_str(self):
        confidence = round(self.confidence, 4) if self.confidence else None
        bbox = " ".join(map(str, self.ltrb))
        return f"{self.class_name}, {bbox}, {confidence}"


class TrackedFrame:
    """A single frame together with its tracked objects."""

    def __init__(
        self,
        frame_id: int,
        frame: numpy.ndarray,
        tracker,
        low_confidence_results: list[list],
        class_names: dict,
    ):
        self.frame_id = frame_id
        self.tracker = copy.deepcopy(tracker)
        self.frame = frame
        self.num_object = len(tracker.tracks)
        self.object_ids = self.get_object_ids()
        self.low_confidence_objects = self.get_low_confidence_objects(
            low_confidence_results, class_names
        )

    def get_low_confidence_objects(
        self, low_confidence_results: list[list], class_names: dict
    ):
        """Returns a list of low confidence objects."""

        low_confidence_objects = []

        for res in low_confidence_results:
            cur_pred = Prediction(
                frame_id=self.frame_id,
                ltrb=ltwh_to_ltrb(res[0]),
                confidence=res[1],
                class_id=res[2],
                class_names=class_names,
            )
            low_confidence_objects.append(cur_pred)

        return low_confidence_objects

    def get_object_ids(self):
        return set(t.track_id for t in self.tracker.tracks)


class TrackedFrameCollection:
    """A collection of TrackedFrames containing frames and the tracked objects."""

    def __init__(
        self,
        video_cap: cv2.VideoCapture,
        class_names: dict,
        out_folder: str,
    ):
        self.video_cap = video_cap
        self.out_folder = out_folder

        self.class_names = class_names
        self.tracked_frames: list = []
        self.all_objects: defaultdict = defaultdict(dict)
        self.all_frames: defaultdict = defaultdict(list)

    def add_tracked_frame(self, tracked_frame: TrackedFrame):
        """Adds a tracked frame to the collection."""

        self.tracked_frames.append(tracked_frame)
        self.update_all_objects_dict(tracked_frame)

    def update_all_objects_dict(self, tracked_frame: TrackedFrame):
        """Updates the dictionary of objects. Each key is an object ID
        and the value is a dictionary of frame IDs and predictions.
        """

        for track in tracked_frame.tracker.tracks:
            frame_id = tracked_frame.frame_id
            cur_pred = Prediction(
                frame_id=tracked_frame.frame_id,
                ltrb=list(map(int, track.to_ltrb())),
                confidence=track.det_conf,
                class_id=track.det_class,
                class_names=self.class_names,
            )
            cur_dict = {tracked_frame.frame_id: cur_pred}
            self.all_objects[track.track_id].update(cur_dict)
            self.all_frames[frame_id].append(cur_pred)

    def export_all_objects(self, out_video_fps: int):
        """Exports all objects to individual videos."""

        os.makedirs(self.out_folder, exist_ok=True)

        for object_id in self.all_objects:
            filename = os.path.join(self.out_folder, f"obj_{object_id}.mp4")
            writer = create_video_writer(
                self.video_cap, filename, fps=out_video_fps
            )
            self.export_object(writer, object_id)

    def export_object(self, writer: cv2.VideoWriter, object_id: str):
        """Exports a single object to a video file."""

        a_dict = self.all_objects[object_id]

        start_idx, end_idx = min(a_dict.keys()), max(a_dict.keys())

        for frame_id in range(start_idx, end_idx + 1):
            if frame_id not in a_dict:
                continue
            frame = self.tracked_frames[frame_id].frame

            black_frame = numpy.zeros_like(frame)
            cur_prediction = a_dict[frame_id]
            x1, y1, x2, y2 = cur_prediction.ltrb
            black_frame[y1 : y2 + 1, x1 : x2 + 1] = frame[
                y1 : y2 + 1, x1 : x2 + 1
            ]
            class_name = cur_prediction.class_name

            text = f"{frame_id=}, {object_id=}, {class_name=}"
            put_text_on_upper_corner(black_frame, text)
            writer.write(black_frame)

        writer.release()

    def get_frame(self, frame_id: int):
        """Returns a frame with the given frame ID."""

        return self.tracked_frames[frame_id].frame

    def get_frame_predictions(self, frame_id: int):
        """Returns the predictions for a single frame."""

        high_confidence_objects = self.all_frames[frame_id]
        low_confidence_objects = self.tracked_frames[
            frame_id
        ].low_confidence_objects
        return high_confidence_objects, low_confidence_objects

temporal_consistency/tracked_frame.py
smttsp/temporal_consistency_odt
import copy
import sys
from datetime import datetime

import cv2


EPS = sys.float_info.epsilon


def create_video_writer(
    video_cap: cv2.VideoCapture, output_filename: str, fps: float = -1
):
    """Create a video writer object to write the output video

    Args:
        video_cap (cv2.VideoCapture): Video capture object
        output_filename (str): Output filename
        fps (float, optional): Frames per second. Defaults to None.

    Returns:
        cv2.VideoWriter: Video writer object
    """

    # grab the width, height, and fps of the frames in the video stream.
    frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(video_cap.get(cv2.CAP_PROP_FPS)) if fps <= 0 else fps

    # initialize the FourCC and a video writer object
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(
        output_filename, fourcc, fps, (frame_width, frame_height)
    )

    return writer


def compute_iou(bbox1: list[int], bbox2: list[int]):
    """Compute the intersection over union (IoU) of two bounding boxes.

    Args:
        bbox1 (list): First bounding box in the format of [x1, y1, x2, y2]
        bbox2 (list): Second bounding box in the format of [x1, y1, x2, y2]

    Returns:
        float: Intersection over union (IoU) of bbox1 and bbox2
    """

    x1 = max(bbox1[0], bbox2[0])
    y1 = max(bbox1[1], bbox2[1])
    x2 = min(bbox1[2], bbox2[2])
    y2 = min(bbox1[3], bbox2[3])

    # Compute the area of the intersection rectangle
    inter_area = max(0, x2 - x1) * max(0, y2 - y1)

    # Compute the area of both bounding boxes
    box1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
    box2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

    # Compute the IoU
    iou = inter_area / float(box1_area + box2_area - inter_area + EPS)

    return iou


def get_runtime_str():
    """Getting datetime as a string

    Returns:
        str: Datetime as a string
    """

    runtime_str = (
        datetime.now()
        .isoformat()
        .replace(":", "")
        .replace("-", "")
        .replace("T", "-")
        .split(".")[0]
    )
    return runtime_str


def ltwh_to_ltrb(ltwh: list[int]):
    """Converts bounding box coordinates from [left, top, width, height] to
    [left, top, right, bottom].
    """

    ltrb = copy.deepcopy(ltwh)
    ltrb[2] += ltrb[0]
    ltrb[3] += ltrb[1]
    return ltrb


def ltrb_to_ltwh(ltrb: list[int]):
    """Converts bounding box coordinates from [left, top, right, bottom] to
    [left, top, width, height].
    """

    ltwh = copy.deepcopy(ltrb)
    ltwh[2] -= ltwh[0]
    ltwh[3] -= ltwh[1]
    return ltwh

temporal_consistency/utils.py
smttsp/temporal_consistency_odt
import cv2


GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
WHITE = (255, 255, 255)


def put_text_on_upper_corner(frame, text):
    """Put text in the upper left corner of the frame."""

    cv2.putText(
        img=frame,
        text=text,
        org=(20, 60),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.9,
        color=WHITE,
        thickness=2,
    )

    return None


def draw_class_name(frame, ltrb_bbox, track_id, class_name, bbox_color=GREEN):
    """Draws the class name and track id on the frame. The class name and track
    id are drawn on top of the bounding box.
    """

    x_min, y_min, x_max, y_max = map(int, ltrb_bbox)
    # draw the bounding box and the track id
    cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), bbox_color, 2)
    cv2.rectangle(
        frame, (x_min, y_min - 20), (x_min + 60, y_min), bbox_color, -1
    )
    new_text = f"{track_id}: {class_name}"
    cv2.putText(
        img=frame,
        text=new_text,
        org=(x_min + 5, y_min - 8),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.5,
        color=WHITE,
        thickness=2,
    )

    return None


def draw_bbox_around_object(frame, track, ltrb_bbox, classes):
    """Draws a bounding box around the object and the class name and track id
    on the frame.
    """

    track_id = track.track_id
    cls = track.det_class
    class_name = classes.get(cls)
    draw_class_name(frame, ltrb_bbox, track_id, class_name)

    return None


def draw_fps_on_frame(frame, total_time):
    """Draws the FPS on the frame."""

    fps = f"FPS: {1000 / total_time:.2f}"
    cv2.putText(
        img=frame,
        text=fps,
        org=(50, 50),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=2,
        color=BLUE,
        thickness=8,
    )

    return None

temporal_consistency/vis_utils.py
smttsp/temporal_consistency_odt
"""PyTest Configurations for Unit-Testing

References:
- https://stackoverflow.com/questions/44377358/how-can-i-display-the-test-name-after-the-test-using-pytest
- https://docs.pytest.org/en/stable/fixture.html
"""

import pytest


@pytest.fixture(
    scope="session",
)
def env_setup(monkeypatch):
    """Setup Environment Variables through PyTest's `monkeypatch` feature."""
    # monkeypatch.setenv(
    #     "GCLOUD_PROJECT",
    #     "samet-project",
    # )
    return None

tests_unit/conftest.py
smttsp/temporal_consistency_odt
import pytest

from temporal_consistency.utils import (
    compute_iou,
    get_runtime_str,
    ltrb_to_ltwh,
    ltwh_to_ltrb,
)


EPS = 1e-10


@pytest.mark.parametrize(
    "box1, box2, expected_iou",
    [
        # Non-overlapping boxes
        ([0, 0, 1, 1], [3, 3, 4, 4], 0.0),
        # Fully overlapping boxes
        ([0, 0, 1, 1], [0, 0, 1, 1], 1.0),
        # Partially overlapping boxes
        ([0, 0, 1, 1], [1, 1, 2, 2], 0.0),
        # Another partially overlapping boxes
        ([0, 0, 2, 2], [1, 1, 3, 3], 1 / 7),
        # Box2 inside box1
        ([0, 0, 3, 3], [1, 1, 2, 2], 1 / 9),
        # Negative coordinates
        ([-3, -3, -1, -1], [-2, -2, -1, -1], 1 / 4),
        # Negative coordinates
        ([-3, -3, 1, 1], [-4, -4, -2, -2], 1 / 19),
    ],
)
def test_compute_iou(box1, box2, expected_iou):
    assert abs(compute_iou(box1, box2) - expected_iou) < EPS


def test_get_runtime_str():
    runtime_str = get_runtime_str()
    assert len(runtime_str) == 15
    assert runtime_str[8] == "-"


@pytest.mark.parametrize(
    "ltrb, ltwh",
    [
        ([1, 2, 3, 4], [1, 2, 2, 2]),
        ([1, 2, 4, 6], [1, 2, 3, 4]),
    ],
)
def test_ltwh_to_ltrb(ltrb, ltwh):
    assert ltwh_to_ltrb(ltwh) == ltrb


@pytest.mark.parametrize(
    "ltrb, ltwh",
    [
        ([1, 2, 8, 6], [1, 2, 7, 4]),
        ([1, 2, 4, 6], [1, 2, 3, 4]),
    ],
)
def test_ltrb_to_ltwh(ltrb, ltwh):
    assert ltrb_to_ltwh(ltrb) == ltwh

tests_unit/test_utils.py
smttsp/temporal_consistency_odt