# this mounts your Google Drive to the Colab VM.
from google.colab import drive
drive.mount('/content/drive', force_remount=True)

dataFolder='data/'
! mkdir '/home/data'
% cd '/home/data'

import os
!pip install --upgrade fastai 
!pip install aicrowd-cli

API_KEY = '52ab6eb031245b7028158e2f3e993174' #Please enter your API Key from [https://www.aicrowd.com/participants/me]
!aicrowd login --api-key $API_KEY
!aicrowd dataset download --challenge f1-car-detection -j 3

!rm -rf data
!mkdir data

!unzip train.zip  -d data/train
!unzip val.zip -d data/val
!unzip test.zip  -d data/test

!mv train.csv data/train.csv
!mv val.csv data/val.csv
!mv sample_submission.csv data/sample_submission.csv

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader,Dataset
import torchvision.transforms as T
import torchvision.models as models


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import ast

import time
import os
import copy

class BoxDataset(Dataset):
  def __init__(self,path,df,imTransf=T.ToTensor(),bbTransf=None):
    self.path=path
    self.df=df
    self.imTransf=imTransf
    self.bbTransf=bbTransf

  def __len__(self):
    return len(self.df)

  def __getitem__(self,ind):
    im=self.load_image(ind)
    bbDots=np.array(ast.literal_eval(self.df.iloc[ind,1]))
    t_im=self.imTransf(im)
    return t_im,bbDots
    
  def load_image(self,ind):
    return Image.open(self.path+str(ind)+'.jpg')

trainTransf=T.Compose([
#        T.Resize(imSize),
#        transforms.RandomHorizontalFlip(),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])

])

df_train=pd.read_csv(dataFolder+'train.csv')
ds_train=BoxDataset(dataFolder+'train/',df_train,trainTransf,bbTransf)
dl_train=DataLoader(ds_train,batch_size=64,shuffle=True,num_workers=2)

df_val=pd.read_csv(dataFolder+'val.csv')
ds_val=BoxDataset(dataFolder+'val/',df_val,trainTransf,bbTransf)
dl_val=DataLoader(ds_val,batch_size=64,num_workers=2)

dataloaders_dict={'val':dl_val,'train':dl_train}

####### training loop
#### code was taken from official pytorch tutorial
def train_model(model, dataloaders, criterion, optimizer, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):

                    outputs = model(inputs)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.detach().item() * inputs.size(0)
                iou=bb_intersection_over_union(labels.detach().cpu().numpy(),outputs.detach().cpu().numpy())
                running_corrects += np.sum(iou >0.5)
                
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects / len(dataloaders[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))


        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

    return model


  #https://www.pyimagesearch.com/2016/11/07/intersection-over-union-iou-for-object-detection/
def bb_intersection_over_union(boxA, boxB):
	# determine the (x, y)-coordinates of the intersection rectangle
	xA = np.maximum(boxA[:,0], boxB[:,0])
	yA = np.maximum(boxA[:,2], boxB[:,2])
	xB = np.minimum(boxA[:,1], boxB[:,1])
	yB = np.minimum(boxA[:,3], boxB[:,3])
	# compute the area of intersection rectangle
	interArea = np.maximum(0, xB - xA + 1) * np.maximum(0, yB - yA + 1)
	# compute the area of both the prediction and ground-truth
	# rectangles
	boxAArea = (boxA[:,1] - boxA[:,0] + 1) * (boxA[:,3] - boxA[:,2] + 1)
	boxBArea = (boxB[:,1] - boxB[:,0] + 1) * (boxB[:,3] - boxB[:,2] + 1)
	# compute the intersection over union by taking the intersection
	# area and dividing it by the sum of prediction + ground-truth
	# areas - the interesection area
	iou = interArea / (boxAArea + boxBArea - interArea)
	# return the intersection over union value
	return iou

!pip install efficientnet_pytorch
from efficientnet_pytorch import EfficientNet

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = EfficientNet.from_pretrained('efficientnet-b3',num_classes = 4)
#model.load_state_dict(torch.load('/content/drive/MyDrive/weights_ef6_net_adam1.txt'))
model.to(device)

num_epochs=5

torch.cuda.empty_cache()
model.to(device)

optimizer =torch.optim.Adam(model.parameters(), lr=0.0001)
criterion = nn.L1Loss(reduction='mean')
model = train_model(model, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs, is_inception=False)

num_epochs=5
optimizer =torch.optim.Adam(model.parameters(), lr=0.0001)
model = train_model(model, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs, is_inception=False)
torch.save(model.state_dict(), '/content/drive/MyDrive/Colab Notebooks/aicrowd/f1/2 challenge/weights_efnet_adam1.txt')

num_epochs=5
optimizer =torch.optim.Adam(model.parameters(), lr=0.00008)
model = train_model(model, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs, is_inception=False)

num_epochs=5
optimizer =torch.optim.Adam(model.parameters(), lr=0.00002)
model = train_model(model, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs, is_inception=False)

num_epochs=4
optimizer =torch.optim.Adam(model.parameters(), lr=0.00002)
model = train_model(model, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs, is_inception=False)

A=[[i for i in range(5000)],['']*5000]
df=pd.DataFrame(A).transpose()
df.columns=['ImageID','bboxes']
i=0
for f in os.listdir('data/test/'):
  im=Image.open('data/test/'+f)
  tens=torch.reshape(trainTransf(im),(1,3,256,256))
  inputs = tens.to(device)
  outputs = (model(inputs).detach().cpu().numpy())
  st='[[' + str(int(outputs[0,0]))+', '+ str(int(outputs[0,2]))+', '+ str(int(outputs[0,1]))+', '+ str(int(outputs[0,3])) +', 0.8]]'
  df.iloc[int(f.split('.')[0]),1]=st


df.to_csv('/content/drive/MyDrive/Colab Notebooks/aicrowd/f1/2 challenge/submission_.csv',index=False)

!aicrowd submission create -c f1-car-detection -f '/content/drive/MyDrive/Colab Notebooks/aicrowd/f1/2 challenge/submission_.csv'