unrotate_annotations.py

import pandas as pd
import numpy as np
import argparse

from bb_tools.obb_to_sbb import get_SBBs


def parse_args():
    parser = argparse.ArgumentParser(description='Convert from oriented bounding boxes in format [cx, cy, w, h, angle] to straight boxes in format [tlx, tly, w, h]')
    
    parser.add_argument('--annot_file', dest='annotation_file', help='Name of the .csv file in MOTChallenge format.')
    parser.add_argument('--out_file', dest='output_file', help='Name of the output .csv file')
    parser.add_argument('--p_length', default=0.15, dest='p_length', help='proportion of the length to remove in the corners.')
    parser.add_argument('--p_width', default=0.15, dest='p_width', help='proportion of the width to remove in the corners.')

    args = parser.parse_args()

    return args


def x1y1wha_to_cxcywha(bbox_x1y1wha):
    x,y,w,h,a = bbox_x1y1wha
    cx, cy = x+w/2, y+h/2
    return np.array([cx,cy,w,h,a])

def cxcywha_to_x1y1wha(bbox_cxcywha):
    cx,cy,w,h,a = bbox_cxcywha
    tlx, tly = cx-w/2, cy-h/2
    return np.array([tlx,tly,w,h,a])


if __name__ == '__main__':

    args = parse_args()

    annotation_file = args.annotation_file
    p_length        = float(args.p_length)
    p_width         = float(args.p_width)
    output_file     = args.output_file


    df = pd.read_csv(annotation_file, header=None)
    df.columns = ['frame', 'id', 'cx', 'cy', 'w', 'h', 'conf', 'x', 'y', 'x', 'angle']

    # Convert to numpy array 
    OBBs = df[['cx','cy', 'w', 'h', 'angle']].to_numpy()

    # Convert to straight bboxes
    SBBs, _ = get_SBBs(OBBs, p_length=0.15, p_width=0.15)

    # Convert to BBoxes to [tlx, tly, w, h, angle] format
    SBBs2 = np.array([[cxcywha_to_x1y1wha(ii) for ii in SBBs]])


    dfs = df.copy()
    dfs = dfs.drop(['angle'], axis=1)
    dfs = dfs.rename(columns={"cx": "tlx", "cy": "tly"})

    dfs[['tlx','tly', 'w', 'h']] = SBBs2[0][:,0:4]

    # Round to 2 decimal places
    dfs = dfs.round({'tlx': 2, 'tly': 2, 'w':2, 'h':2})

    print (dfs)
    dfs.to_csv(output_file, index=False, header=None)