detector2cvat.py

import os
from pathlib import Path
import json
import argparse
import cv2
from tqdm import tqdm
from kabr_tools.utils.yolo import YOLOv8
from kabr_tools.utils.tracker import Tracker, Tracks
from kabr_tools.utils.object import Object
from kabr_tools.utils.draw import Draw


def detector2cvat(path_to_videos: str, path_to_save: str,
                  target_labels: list, label_map: dict, model: str, show: bool) -> None:
    """
    Detect objects with Ultralytics YOLO detections, apply SORT tracking and convert tracks to CVAT format.

    Parameters:
    path_to_videos - str. Path to the folder containing videos.
    path_to_save - str. Path to the folder to save output xml & mp4 files.
    target_labels - list. List of target labels to detect.
    label_map - dict. Dictionary to rename labels.
    model - str. YOLO model to use with detections.
    show - bool. Flag to display detector's visualization.
    """
    videos = []

    for root, dirs, files in os.walk(path_to_videos):
        for file in files:
            if os.path.splitext(file)[1] == ".mp4":
                folder = Path(root).name

                if folder.startswith("!") or file.startswith("!"):
                    continue

                videos.append(str(Path(root) / file))

    yolo = YOLOv8(weights=model, imgsz=3840, conf=0.5, target_labels=target_labels, label_map=label_map)

    for i, video in enumerate(videos):
        try:
            name = Path(video).stem

            video_parts = [p for p in Path(video).parts[-3:-1] if not Path(p).is_absolute()]
            output_folder = Path(path_to_save).joinpath(*video_parts)
            output_path = output_folder / f"{name}.xml"
            print(f"{i + 1}/{len(videos)}: {video} -> {output_path}")

            if not output_folder.exists():
                output_folder.mkdir(parents=True)

            vc = cv2.VideoCapture(video)
            size = int(vc.get(cv2.CAP_PROP_FRAME_COUNT))
            width = int(vc.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(vc.get(cv2.CAP_PROP_FRAME_HEIGHT))
            vw = cv2.VideoWriter(str(output_folder / f"{name}_demo.mp4"), cv2.VideoWriter_fourcc("m", "p", "4", "v"),
                                 29.97, (width, height))
            max_disappeared = 40
            tracker = Tracker(max_disappeared=max_disappeared, max_distance=300)
            tracks = Tracks(max_disappeared=max_disappeared, interpolation=True,
                            video_name=name, video_size=size, video_width=width, video_height=height)
            index = 0
            vc.set(cv2.CAP_PROP_POS_FRAMES, index)
            pbar = tqdm(total=size)

            while vc.isOpened():
                returned, frame = vc.read()

                if returned:
                    visualization = frame.copy()
                    predictions = yolo.forward(frame)
                    centroids = []
                    attributes = []

                    for prediction in predictions:
                        attribute = {}
                        centroids.append(YOLOv8.get_centroid(prediction[0]))
                        attribute["box"] = prediction[0]
                        attribute["confidence"] = prediction[1]
                        attribute["label"] = prediction[2]
                        attributes.append(attribute)

                    objects, colors = tracker.update(centroids)
                    objects = Object.object_factory(objects, centroids, colors, attributes=attributes)
                    tracks.update(objects, index)

                    for object in objects:
                        Draw.track(visualization, tracks[object.object_id].centroids, object.color, 20)
                        Draw.bounding_box(visualization, object)
                        Draw.object_id(visualization, object)

                    cv2.putText(visualization, f"Frame: {index}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX,
                                0.8, (255, 255, 255), 3, cv2.LINE_AA)
                    if show:
                        cv2.imshow("detector2cvat", cv2.resize(
                            visualization, (int(width // 2.5), int(height // 2.5))))
                    vw.write(visualization)
                    key = cv2.waitKey(1)
                    index += 1
                    pbar.update(1)

                    if key == 27:
                        break
                else:
                    break

            pbar.close()
            vc.release()
            vw.release()
            cv2.destroyAllWindows()
            tracks.save(str(output_path), "cvat")
        except Exception as e:
            print(f"Something went wrong: {e}")


def parse_args() -> argparse.Namespace:
    local_parser = argparse.ArgumentParser()
    local_parser.add_argument(
        "--video",
        type=str,
        help="path to folder containing videos",
        required=True
    )
    local_parser.add_argument(
        "--save",
        type=str,
        help="path to save output xml & mp4 files",
        required=True
    )
    local_parser.add_argument(
        "--target_labels",
        type=str,
        help="path to target labels json"
    )
    local_parser.add_argument(
        "--label_map",
        type=str,
        help="path to label map json"
    )
    local_parser.add_argument(
        "--yolo",
        type=str,
        default="yolov8x.pt",
        help="yolo model to use with detections"
    )
    local_parser.add_argument(
        "--imshow",
        action="store_true",
        help="flag to display detector's visualization"
    )
    return local_parser.parse_args()

def load_json(file: str) -> dict:
    if file:
        with open(file, mode="r", encoding="utf-8") as file:
            return json.load(file)
    return None

def main() -> None:
    args = parse_args()
    target_labels = load_json(args.target_labels)
    label_map = load_json(args.label_map)
    detector2cvat(args.video, args.save, target_labels, label_map, args.yolo, args.imshow)


if __name__ == "__main__":
    main()