!pip install git+https://gitlab.aicrowd.com/yoogottamk/aicrowd-cli.git
from aicrowd.magic import *
# Get your API key from https://www.aicrowd.com/participants/me
%aicrowd login --api-key ""

%aicrowd dataset download --challenge objde *v2.*

!rm -rf data
!mkdir data

!unzip train-v2.zip -d data/

!unzip val-v2.zip -d data/

!unzip test-v2.zip -d data/

!mv train-v2.csv data/train-v2.csv
!mv val-v2.csv data/val-v2.csv

# install dependencies: 
!pip install pyyaml==5.1
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.7/index.html

# Some basic setup:
# Setup detectron2 logger
import torch, torchvision

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, cv2, random
from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.structures import BoxMode

import pandas as pd
import numpy as np
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
import cv2

train_df = pd.read_csv("data/train-v2.csv")
val_df = pd.read_csv("data/val-v2.csv")
train_df

def show_images(images, num = 5):
    
    images_to_show = np.random.choice(images, num)

    for image_id in images_to_show:

        image = Image.open(f'data/train-v2/{image_id}.jpg')

        w, h = image.size

        print(w, h)

        bboxes = train_df[train_df['ImageID'] == image_id][['XMin','XMax','YMin','YMax']].to_numpy()

        # visualize them
        draw = ImageDraw.Draw(image)
        for bbox in bboxes:    
            draw.rectangle([bbox[0]*w, bbox[2]*h, bbox[1]*w, bbox[3]*h], width=3)

        plt.figure(figsize = (15,15))
        plt.imshow(image)
        plt.show()

show_images(train_df['ImageID'].unique(), num = 1)

category_dict = {"Person":0, "Clothing":1, "Car":2, "Plant":3, "Footwear":4}

train_df['category_id'] = train_df['LabelName'].apply(lambda x : category_dict[x])
train_df['image_category_id'] = train_df.groupby(['ImageID']).ngroup() 
train_df

for i, row in train_df[train_df["ImageID"] == "0"].iterrows():

  print("row", row["class_id"])

dict_dataset = []
def get_dataset_dics(img_dir):

    for i in tqdm(train_df.groupby('ImageID')):

        image = Image.open(f'data/train-v2/{i[0]}.jpg')
        w, h = image.size
        
        ann_lst = []

        bboxes = i[1][['XMin','XMax','YMin','YMax']].to_numpy()

        for n, bbox in enumerate(bboxes):   
    
            ann_dict = {'bbox': [bbox[0]*w, bbox[2]*h, bbox[1]*w, bbox[3]*h],
           'bbox_mode': BoxMode.XYXY_ABS,
           'category_id': i[1]['category_id'].values[0],
           'iscrowd': 0}
            
            ann_lst.append(ann_dict)

        image_dict = {'annotations': ann_lst,
            'file_name': 'data/train-v2/'+i[0]+'.jpg',
            'height': h,
            'image_id': i[1]['image_category_id'].values[0],
            'width': w}
          
        dict_dataset.append(image_dict)

    return dict_dataset

dict_dataset = get_dataset_dics(train_df['ImageID'])

from detectron2.data import DatasetCatalog, MetadataCatalog

d = f"obj_train{np.random.randint(10000)}"
DatasetCatalog.register(d, lambda d=d : get_dataset_dics(train_df['ImageID']))
MetadataCatalog.get(d).set(thing_classes=["Person", "Clothing", "Car", "Plant", "Footwear"])
obj_metadata = MetadataCatalog.get(d)

for i in random.sample(dict_dataset, 3):
    img = cv2.imread(i["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=obj_metadata, scale=0.5)
    out = visualizer.draw_dataset_dict(i)
    cv2_imshow(out.get_image()[:, :, ::-1])

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_DC5_3x.yaml"))
cfg.DATASETS.TRAIN = (d,)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_DC5_3x.yaml") 
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 200
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128  
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

%load_ext tensorboard
%tensorboard --logdir output

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 
predictor = DefaultPredictor(cfg)

test_imgs_paths = os.listdir("data/test-v2")

full_predictions = {"ImageID":[], "XMin":[], "XMax":[], "YMin":[], "YMax":[], "score":[], "LabelName":[]}

reverse_category_dict = {0:"Person", 1:"Clothing", 2:"Car", 3:"Plant", 4:"Footwear"}

for test_img_path in tqdm(test_imgs_paths):

  img = cv2.imread(f"data/test-v2/{test_img_path}")
  h, w, _ = img.shape

  predictions = predictor(img)

  bboxes = predictions['instances'].pred_boxes.tensor.cpu().numpy()

  new_boxes = []

  for b in bboxes:

    new_boxes.append([b[0]*w, b[2]*h, b[1]*w, b[3]*h])

  if new_boxes != []:

    new_boxes = np.array(new_boxes)

    classes = predictions['instances'].pred_classes.cpu().numpy()
    scores = predictions['instances'].scores.cpu().numpy()
    image_id = [test_img_path.split('.')[0]]*scores.shape[0]

    classes = [reverse_category_dict[i] for i in classes]

    full_predictions['ImageID'].extend(image_id)
    full_predictions['XMin'].extend(new_boxes[:, 0])
    full_predictions['XMax'].extend(new_boxes[:, 1])
    full_predictions['YMin'].extend(new_boxes[:, 2])
    full_predictions['YMax'].extend(new_boxes[:, 3])
    full_predictions['score'].extend(scores)
    full_predictions['LabelName'].extend(classes)
  
  else:
    pass

submission = pd.DataFrame(full_predictions)
submission

submission.to_csv("submission.csv", index=False)

try:
    from google.colab import files
    files.download('submission.csv') 
except:
    print("Option Only avilable in Google Colab")