# default_exp datasets.generators

Generators for MOT data¶

MOT data format can be found here.

from pandas import DataFrame
import numpy as np
from pathlib import Path
from skimage.draw import rectangle

import tempfile
tmp_dir = tempfile.TemporaryDirectory()
tmp_dir.name

'/tmp/tmpmh6z1wli'

def show_img(img):
    fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(3, 3))
    ax.imshow(img)

Create mot gt and render corresponding frames¶

#export
def create_mot_ds(path: Path, img_path: str = 'img', n_frames: int = 20, occluded: bool = False):
    # create data
    frames = gen_frame(n_frames=n_frames)
    track1 = gen_moving_object(frames, object_id=0, size=(10, 30),
                               path_f=partial(default_path, start=(150, 10), speed=(0, 10)))
    track2 = gen_moving_object(track1, object_id=1, type='circle', color=(1, 0, 0), size=(40, 40),
                               path_f=partial(default_path, start=(40, 40), speed=(5, 5)))
    bg_img = gen_bg_img(track2)
    track_img = gen_render_objects(bg_img)

    if occluded:
        track_img = occlude_lower_left_img(track_img)

    # save to files
    _img_path = Path(path) / img_path
    _img_path.mkdir(parents=True, exist_ok=True)
    records = []

    for gt in track_img:
        frame_pos = str(gt.pos).zfill(4)
        file_path = _img_path / f"{frame_pos}.jpg"
        im = Image.fromarray((gt.img * 175.5).astype(np.uint8))  # Trasnform img to RGB

        im.save(file_path)

        for mot in gt_to_mot(gt):
            records.append(mot)

    df_gt = DataFrame.from_records(records)
    df_gt.to_csv(Path(path) / 'mot.csv', index=False)

Dataset factory

Navi Widget