A parser for labeled fishes in the wild

lfiw.py

import json
import re
from typing import List

from dive_utils.models import Feature, Track


def load_dat(path: str, outpath: str):
    with open(path, "r") as f:
        with open(outpath, "w") as f2:
            tracks = load_dat_as_tracks_and_attributes(f.readlines())
            out = {}
            for track in tracks:
                out[track.trackId] = track.dict()
            f2.write(json.dumps(out))


def load_dat_as_tracks_and_attributes(rows: List[str]) -> List[Track]:
    """
    The original Labeled fishes in the wild dataset (v1.0, Dec. 2014)
    contained only the decimated test video sequence ("Test_ROV_video_h264_decim.mp4")
    that contained only the marked frames from the original video. One tenth of the frames
    of the full frame-rate video were marked for locations of fish targets. This version
    of the dataset (v1.1) also contains the full test video sequence
    ("Test_ROV_video_h264_full.mp4").
    Both the full and decimated videos have accompanying text files with analyst marks
    (following OpenCV .dat file conventions).

    Generally, for m marks, the format is:

    Video-filename(frame#) #-of-marks x1 y1 w1 h1 x2 y2 w2 h2 ... xm ym wm hm

    For example, in the case of two marks, the final eight values define the bounding rectangles:
    Test_ROV_video_h264_full.mp4(fr_14) 2 1021 362 94 63 953 289 90 61

    The marks file for the decimated video ("Test_ROV_video_h264_decim_marks.dat") indicates the
        frame number for the decimated and full sequence.

    Test_ROV_video_h264_decim.mp4(fr_1)(fullfr_14) 2 1021 362 94 63 953 289 90 61

    There are 2101 frames in the full video and 210 frames in the decimated video,
        but 206 frames were marked; i.e. a few frames did not contain verified fish.
    """
    tracks = []
    trackId = 0
    for row in rows:
        columns = row.split(' ')
        print(row, columns)
        frame = re.search(r'\(fr_(\d*)\)', columns[0]).group(1)
        frame = int(frame) - 1
        marks = int(columns[1])
        for j in range(marks):
            offset = 2
            jump = 4
            x = columns[offset + (j * jump)]
            y = columns[offset + (j * jump) + 1]
            w = columns[offset + (j * jump) + 2]
            h = columns[offset + (j * jump) + 3]
            bounds = [
                int(x),
                int(y),
                int(x) + int(w),
                int(y) + int(h),
            ]
            feature = Feature(bounds=bounds, frame=frame)
            tracks.append(
                Track(
                    begin=frame,
                    end=frame,
                    trackId=trackId,
                    features=[feature],
                    confidencePairs=[('fish', 1.0)],
                )
            )
            trackId += 1
    return tracks