MOTDet17 to Yolov5 (people class)

mot17det_to_yolo.py

# wget https://motchallenge.net/data/MOT17Det.zip
# unzip MOT17Det.zip
# rm -rf MOT17Det/test/

import configparser
import os 
import shutil
import numpy as np

def get_metadata_from_seqinfo(seqinfo_path):
    print(seqinfo_path)
    config = configparser.ConfigParser()
    config.read(seqinfo_path)
    width, height = config['Sequence']['imWidth'], config['Sequence']['imHeight']
    im_dir, imExt, fps, length = config['Sequence']['imDir'], config['Sequence']['imExt'], config['Sequence']['frameRate'], config['Sequence']['seqLength']
    return {"width" : int(width), "height" : int(height), "im_dir" : im_dir, "imExt" : imExt,"fps" : int(fps), "length" : int(length)}


def parse_line(annotation_line):
    nums = [x for x in annotation_line.split(',')]
    bbox = nums[2:6] # left, top, width, height
    label = int(nums[7])
    frame = int(nums[0])
    return frame, label, bbox

def bbox_conversion(bbox_mot, img_size):
    # from list to str with numbers
    width, height = img_size
    bbox_tl_x = float(bbox_mot[0]) / float(width)
    bbox_tl_y = float(bbox_mot[1]) / float(height)
    bbox_w = float(bbox_mot[2]) / float(width)
    bbox_h = float(bbox_mot[3]) / float(height)

    bbox_tl_x = np.clip(bbox_tl_x, 0, 1)
    bbox_tl_y = np.clip(bbox_tl_y, 0, 1)
    bbox_w = np.clip(bbox_w, 0, 1)
    bbox_h = np.clip(bbox_h, 0, 1)

    bbox_center_x  = bbox_tl_x + bbox_w*0.5
    bbox_center_y = bbox_tl_y + bbox_h*0.5

    bbox_center_x = np.clip(bbox_center_x, 0, 1)
    bbox_center_y = np.clip(bbox_center_y, 0, 1)

    bbox_str = "{} {} {} {}".format(bbox_center_x, bbox_center_y, bbox_w, bbox_h)
    return bbox_str

def parse_mot_gt(gt_path):
    gt_data={}
    annotation_file = open(gt_path)
    gt_lines = annotation_file.readlines()
    for line in gt_lines:
        frame, label, bbox = parse_line(line)
        if frame not in gt_data:
            gt_data[frame] = []
        gt_data[frame].append((label, bbox))
    return gt_data

# parsing train folder:
MOTdet_root="MOT17Det/train/"
MOTdet_output_images_yolo="MOT17Det/images"
MOTdet_output_annotations_yolo="MOT17Det/labels"

if not os.path.exists(MOTdet_output_images_yolo):
    os.makedirs(MOTdet_output_images_yolo)
if not os.path.exists(MOTdet_output_annotations_yolo):
    os.makedirs(MOTdet_output_annotations_yolo)

pedestrian_cls = 1
person_on_vehicle = 2
static_person = 7
distractor_id = 8
reflection_id = 12
person_classes = [pedestrian_cls, person_on_vehicle, static_person]
label_list=person_classes

seq_folders= os.listdir(MOTdet_root)
for seq_folder in seq_folders:
    print(seq_folder)
    seq_path=os.path.join(MOTdet_root, seq_folder)
    seq_meta = get_metadata_from_seqinfo(os.path.join(seq_path, "seqinfo.ini"))
    print(seq_meta)

    gt_path = os.path.join(os.path.join(seq_path, "gt"), "gt.txt")
    gt_dict = parse_mot_gt(gt_path)

    img_size = (seq_meta['width'], seq_meta['height'])

    # handle images
    for i in range(1, int(seq_meta['length']+1)):
        img_name=('{:06d}'.format(i)) + seq_meta['imExt']
        source = os.path.join(os.path.join(seq_path,seq_meta['im_dir']), img_name)
        img_name_dst = seq_folder + '_' + img_name
        destination = os.path.join(MOTdet_output_images_yolo, img_name_dst)
        shutil.copy2(source, destination)

        # Annotations
        annotation_name = img_name_dst.replace(seq_meta['imExt'],"") + ".txt"
        annotation_name_path = os.path.join(MOTdet_output_annotations_yolo, annotation_name)
        print(annotation_name_path)
        frame_annotations = gt_dict[i]
        with open(annotation_name_path, 'w') as annotation_file:
            for frame_annotation in frame_annotations:
                cur_label = frame_annotation[0]
                if cur_label in label_list:
                    cur_label=0 # 1 class: Person
                    annotation_file.write("{} {}\n".format(cur_label, bbox_conversion(frame_annotation[1], img_size)))