import json
import random
random.seed(2021)
import os, sys
from IPython.display import display, clear_output, HTML
from random import randrange
import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"]=25,25
import numpy as np
import seaborn as sns

# Because Life, Universe and Everything!
random.seed(42)

def mdprint(text):
    display({
        'text/markdown': text,
        'text/plain': text
    }, raw=True)

%cd /content
!git clone http://gitlab.aicrowd.com/amazon-prime-air/airborne-detection-starter-kit.git
os.chdir("airborne-detection-starter-kit/data")

# Dataset for Airborne Object Tracking Dataset
sys.path.append(os.path.dirname(os.path.realpath(os.getcwd())))
sys.path.append(os.path.dirname(os.path.realpath(os.getcwd())) + "/core")
!pip install -r ../requirements.txt > /dev/null
from core.dataset import Dataset
notebook_path = os.path.dirname(os.path.realpath("__file__"))

local_path = notebook_path + '/part1'
s3_path = 's3://airborne-obj-detection-challenge-training/part1/'
dataset = Dataset(local_path, s3_path)

# keys
def remove_numbers(s):
  return ''.join([i for i in s if not i.isdigit()])

all_keys = []
for flight_id in dataset.get_flight_ids():
    flight = dataset.get_flight(flight_id)
    all_keys.extend([remove_numbers(k) for k in flight.detected_objects])

unique_keys = list(set(all_keys))
unique_keys

rows = []
for flight_id in dataset.get_flight_ids():
    flight = dataset.get_flight(flight_id)
    for obj_key in flight.detected_objects:
        object_type = remove_numbers(obj_key)
        obj = flight.detected_objects[obj_key]
        for loc in obj.location:
            bbox = loc.bb.get_bbox()
            frame_id = loc.frame.id
            range_distance = loc.range_distance_m
            image_path = loc.frame.image_path()
            rows.append([flight_id, object_type, obj_key, frame_id, 
                         *bbox, bbox[-1]*bbox[-2], image_path, range_distance])

import pandas as pd
df = pd.DataFrame(rows)
df.columns = ['flight_id', 'object_type', 'object', 'frame_id', 
              'left', 'top', 'width', 'height', 'area', 'image_path','range_distance']
df.head()

import random
filtered_df = df.loc[((df['object_type'] == 'Helicopter') | (df['object_type'] == 'Airplane'))]
filtered_df = filtered_df.loc[df['area'] > 400]
flights = filtered_df.index[(filtered_df['object_type'] == 'Airplane')]
helicopters = filtered_df.index[(filtered_df['object_type'] == 'Helicopter')]
counts_dict = dict(filtered_df['image_path'].value_counts())
multi_object_frames = filtered_df[filtered_df.apply(lambda x:counts_dict[x['image_path']] > 1, axis=1)].index
train_indices = random.sample(list(flights),5000) + random.sample(list(helicopters),5000)
train_df = filtered_df.loc[list(set(train_indices))]
filtered_df.head()

download_images = []
for idx in range(len(train_df)):
  flight = dataset.flights[train_df.iloc[idx]['flight_id']]
  frame = flight.frames[train_df.iloc[idx]['frame_id']]
  image_path = frame.image_path()
  s3_path = frame.image_s3_path()
  download_images.append([image_path, image_path])
random.shuffle(download_images)

from tqdm import tqdm
for idx in tqdm(range(100)):
  dataset.file_handler.download_from_s3_parallel(download_images[idx])

import cv2
import sys
%matplotlib inline

def drawBoundingBoxes(imageInputPath, imageOutputPath, inferenceResults, color):
    """Draw bounding boxes on an image.
    imageData: image data in numpy array format
    imageOutputPath: output image file path
    inferenceResults: inference results array off object (l,t,w,h)
    colorMap: Bounding box color candidates, list of RGB tuples.
    """
    imageData = cv2.imread(imageInputPath)
    for res in inferenceResults:
        left = int(res['left'])
        top = int(res['top'])
        right = int(res['left']) + int(res['width'])
        bottom = int(res['top']) + int(res['height'])
        label = res['label']
        imgHeight, imgWidth, _ = imageData.shape
        thick = int((imgHeight + imgWidth) // 900)
        print (left, top, right, bottom)
        cv2.rectangle(imageData,(left, top), (right, bottom), color, thick)
        cv2.putText(imageData, label, (left, top - 12), 0, 1e-3 * imgHeight, color, thick//3)
    plt.imshow(imageData)

target_image = download_images[0][0]
color = (0,255,0)
img_root = 'part1'
frame_df = train_df[df["image_path"] == target_image]
boxes = []
for idx,ob in frame_df.iterrows():
  label,left,top,width,height = ob['object_type'],ob['left'],ob['top'],ob['width'],ob['height']
  box = {}
  box['left'] = left
  box['top'] = top
  box['width'] = width
  box['height'] = height
  box['label'] = label
  boxes.append(box)

drawBoundingBoxes(os.path.join(img_root,target_image),'output.png',boxes,color)

%cd /content/
!git clone https://github.com/AlexeyAB/darknet

%cd darknet/
!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137
!cp yolov4.conv.137 build/darknet/x64

!mkdir data/obj
!cp -r /content/airborne-detection-starter-kit/data/part1/Images data/obj/

import os
import random 
path ="data/obj/"
filelist = []

for root, dirs, files in os.walk(path):
	for file in files:
		filelist.append(os.path.join(root,file))

random.shuffle(filelist)
train_list = filelist[:int(0.8*len(filelist))]
val_list = filelist[int(0.8*len(filelist)):]             

train_file = open('data/obj/train.txt','w')
val_file = open('data/obj/test.txt','w')
object_label_map = {'Airplane':0,'Helicopter':1}
labels_path = "data/obj"

for image_file in train_list:
  train_file.write(image_file+"\n")
  image_key = image_file.replace(path,"")
train_file.close()

for image_file in val_list:
  val_file.write(image_file+"\n")
  image_key = image_file.replace(path,"")
val_file.close()

img_width = 2448
img_height = 2048

for image_file in filelist:
  image_key = image_file.replace(path,"")
  frame_df = train_df[df["image_path"] == image_key]
  label_file = os.path.join(labels_path,image_key.replace("png","txt"))
  os.makedirs(os.path.dirname(label_file), exist_ok=True)
  label_fp = open(label_file,"w")
  for idx,ob in frame_df.iterrows():
    label,left,top,width,height = ob['object_type'],ob['left'],ob['top'],ob['width'],ob['height']
    x_center = left + width/2
    y_center = top + height/2
    label = object_label_map[label]
    label_fp.write(f"{label} {x_center/img_width} {y_center/img_height} {width/img_width} {height/img_height}")
    label_fp.write("\n")
  label_fp.close()

!make