!pip install aicrowd-cli
%load_ext aicrowd.magic

%aicrowd login

!rm -rf data
!mkdir data
%aicrowd ds dl -c mask-prediction -o data

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

!git clone https://github.com/ultralytics/yolov5  # clone
%cd yolov5
!pip install -r requirements.txt

import pandas as pd
import numpy as np
import os

name2id = {"surgical": 0, "N95": 1, "KN95": 2, "cloth": 3}
id2name = ["surgical", "N95",    "KN95", "cloth"]

!mkdir "../data/train/images"
!mkdir "../data/train/labels"
!mkdir "../data/val/images"
!mkdir "../data/val/labels"
!mv ../data/train/*.jpg "../data/train/images"
!mv ../data/val/*.jpg "../data/val/images"

!echo 'train: ../data/train/images' > masks.yaml
!echo 'val: ../data/val/images' >> masks.yaml
!echo '' >> masks.yaml
!echo 'nc: 4' >> masks.yaml
!echo '' >> masks.yaml
!echo 'names: ["surgical", "N95", "KN95", "cloth"]' >> masks.yaml

for folder in ('train', 'val'):
    labels = pd.read_csv(f'../data/{folder}.csv')
    for name, bbox, masktype in zip(labels['ImageID'], labels['bbox'], labels['masktype']):
        bbox = [int(x)/512 for x in bbox[1:-1].split(",")]
        xmin, ymin, xmax, ymax = bbox
        bbox = [(xmax+xmin)/2, (ymax+ymin)/2, xmax-xmin, ymax-ymin]
        bbox = [str(x) for x in bbox]
        masktype = name2id[masktype]

        f = open(f"../data/{folder}/labels/{name}.txt", "w+")
        f.write(f"{masktype} {' '.join(bbox)}")
        f.close()

#optional
%pip install -q wandb
import wandb
wandb.login()

!python train.py --img 512 --batch 25 --epochs 6 --data masks.yaml --weights yolov5l.pt --cache --optimizer=Adam

!python detect.py --imgsz=512 --conf-thres=0 --iou-thres=0.5 --max-det=1 --save-txt --nosave --weights runs/train/exp/weights/best.pt --source ../data/test

submission = pd.read_csv('../data/sample_submission.csv')

for i, row in enumerate(submission['ImageID']):
    x = open(f"runs/detect/exp/labels/{row}.txt").read().split()
    mtype = id2name[int(x[0])]
    x_center, y_center, width, height = [float(i)*512 for i in x[1:]]
    xmin, ymin, xmax, ymax = round(x_center - width/2), round(y_center - height/2), round(x_center + width/2), round(y_center + height/2)
    bbox = f"({xmin},{ymin},{xmax},{ymax})"
    submission.loc[i, 'bbox'] = bbox
    submission.loc[i, 'masktype'] = mtype

submission.head()

!rm -rf assets
!mkdir assets
submission.to_csv(os.path.join("assets", "submission.csv"))

%aicrowd notebook submit -c mask-prediction -a assets --no-verify