# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import json
import pandas as pd
import os
import glob
import re
from tqdm import tqdm_notebook as tqdm
from PIL import Image
import shutil

# +
DS_FOLDER = '/APL/Datasets/UAVBenchmark/'
GT_FOLDER = os.path.join(DS_FOLDER, 'UAV-benchmark-MOTD_v1.0','GT')
IMG_FOLDER = os.path.join(DS_FOLDER, 'UAV-benchmark-M')
ATTR_FOLDER_TRAIN = os.path.join(DS_FOLDER, 'M_attr', 'train')
ATTR_FOLDER_TEST = os.path.join(DS_FOLDER, 'M_attr', 'test')

OUTPUT_DIR = '/APL/Datasets/UAVBenchmark-COCO-format/'
# -

ground_truth_files = glob.glob(os.path.join(GT_FOLDER, '*_gt_whole.txt'))

train_attribute_files = glob.glob(os.path.join(ATTR_FOLDER_TRAIN, '*.txt'))
test_attribute_files = glob.glob(os.path.join(ATTR_FOLDER_TEST, '*.txt'))

id_getter = re.compile('(?:train|test)\/M([0-9]*)_attr')

train_ids = [id_getter.search(i)[1] for i in train_attribute_files]
test_ids = [id_getter.search(i)[1] for i in test_attribute_files]

train_ids

test_ids

#SET = 'train'
SET = 'test'

new_data_format = {'info': {'description': 'UAVDT Benchmark',
  'url': 'https://sites.google.com/site/daviddo0323/projects/uavdt',
  'version': '1.0',
  'year': 2019,
  'contributor': '<NAME> et. al., this format: <NAME> @McGill University',
  'date_created': '2019/07/07'},
                  'images': [],
                  'annotations': [],
                  'categories': []}

# +
# Get all image data
if SET == 'train':
    id_list = train_ids
elif SET == 'test':
    id_list = test_ids
    
image_data = []
    
regex = re.compile('img([0-9]*).')

# Prepare folder to copy images to
IMG_DEST_FOLDER = os.path.join(OUTPUT_DIR, SET)
try:
    os.mkdir(IMG_DEST_FOLDER)
except FileExistsError:
    shutil.rmtree(IMG_DEST_FOLDER)
    os.mkdir(IMG_DEST_FOLDER)
    
# Create list of all image files
for ID in tqdm(id_list):
    im_files = glob.glob(os.path.join(IMG_FOLDER, 'M' + ID, '*.jpg'))
    
    for im_file in tqdm(im_files, leave=False):
        im = Image.open(im_file)
        im_size = im.size
        im.close()

        sequence_num = regex.search(os.path.basename(im_file))[1]
        
        new_id = ID + sequence_num
        new_file_name = 'img' + new_id + '.jpg'
        
        shutil.copy2(im_file, os.path.join(IMG_DEST_FOLDER, new_file_name))

        new_data_format['images'].append(
            {'file_name': new_file_name,
             'height': im_size[1],
             'width': im_size[0],
             'id': int(new_id)})
# -

# Fill out categories
new_data_format['categories'] = [{'id': 1, 'name': 'car', 'supercategory': None}, {'id': 2, 'name': 'truck', 'supercategory': None}, {'id': 3, 'name': 'bus', 'supercategory': None}]

# +
id_getter_2 = re.compile('GT\/M([0-9]*)_gt')

annot_id = 0

# Get annotations
for file in tqdm(ground_truth_files):
    set_id = id_getter_2.search(file)[1]
    data = pd.read_csv(file, header=None, names=['frame_index','target_id','bbox_left','bbox_top','bbox_width','bbox_height','out-of-view','occlusion','object_category'])
    
    for row in tqdm(data.itertuples(), leave=False):
        
        findex = '{:06d}'.format(row.frame_index)
        
        new_data_format['annotations'].append(
            {'id': annot_id,
             'image_id': int(set_id + findex),
             'category_id': row.object_category,
             'bbox': [row.bbox_left, row.bbox_top, row.bbox_width, row.bbox_height],
             'segmentation': [[row.bbox_left, row.bbox_top, row.bbox_left, row.bbox_top + row.bbox_height, row.bbox_left + row.bbox_width, row.bbox_top + row.bbox_height, row.bbox_left + row.bbox_width, row.bbox_top]],
             'area': row.bbox_width * row.bbox_height,
             'iscrowd': 0})
        annot_id += 1
# -

new_data_format['annotations']

new_data_format['images']

OUTPUT_JSON = os.path.join(OUTPUT_DIR, SET + '.json')
#Output to file
with open(OUTPUT_JSON, 'w') as f:
    json.dump(new_data_format,f)



StarcoderdataJupyterScriptsDedupFiltered