"""
## Dataset: [Face Mask Detection](https://www.kaggle.com/andrewmvd/face-mask-detection)
"""
!git clone https://github.com/rkuo2000/yolov5
%cd yolov5
"""
## Repro [YOLOv5](https://github.com/ultralytics/yolov5)
"""
!mkdir -p Dataset/FaceMask/Images
!mkdir -p Dataset/FaceMask/Labels
# copy image files
!cp -rf /kaggle/input/face-mask-detection/images/* Dataset/FaceMask/Images
!mkdir -p Dataset/images Dataset/labels
"""
## Create Dataset
"""
import os
import numpy as np
from pathlib import Path
from xml.dom.minidom import parse
from shutil import copyfile
FILE_ROOT = "/kaggle/input/face-mask-detection/"
IMAGE_PATH = FILE_ROOT + "images"  
ANNOTATIONS_PATH = FILE_ROOT + "annotations"

DATA_ROOT = "Dataset/"
LABELS_ROOT = DATA_ROOT + "FaceMask/Labels"
IMAGES_ROOT = DATA_ROOT + "FaceMask/Images"  

DEST_IMAGES_PATH = "images"
DEST_LABELS_PATH = "labels" 
classes = ['with_mask', 'without_mask', 'mask_weared_incorrect']
"""
### convert annotations (from COCO .xml to YOLO format .txt)
"""
def cord_converter(size, box):
    """
    convert xml annotation to darknet format coordinates
    :param size： [w,h]
    :param box: anchor box coordinates [upper-left x,uppler-left y,lower-right x, lower-right y]
    :return: converted [x,y,w,h]
    """
    x1 = int(box[0])
    y1 = int(box[1])
    x2 = int(box[2])
    y2 = int(box[3])

    dw = np.float32(1. / int(size[0]))
    dh = np.float32(1. / int(size[1]))

    w = x2 - x1
    h = y2 - y1
    x = x1 + (w / 2)
    y = y1 + (h / 2)

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return [x, y, w, h]

def save_file(img_jpg_file_name, size, img_box):
    save_file_name = LABELS_ROOT + '/' + img_jpg_file_name + '.txt'
    print(save_file_name)
    file_path = open(save_file_name, "a+")
    for box in img_box:

        cls_num = classes.index(box[0])

        new_box = cord_converter(size, box[1:])

        file_path.write(f"{cls_num} {new_box[0]} {new_box[1]} {new_box[2]} {new_box[3]}\n")

    file_path.flush()
    file_path.close()
    
def get_xml_data(file_path, img_xml_file):
    img_path = file_path + '/' + img_xml_file + '.xml'
    print(img_path)

    dom = parse(img_path)
    root = dom.documentElement
    img_name = root.getElementsByTagName("filename")[0].childNodes[0].data
    img_size = root.getElementsByTagName("size")[0]
    objects = root.getElementsByTagName("object")
    img_w = img_size.getElementsByTagName("width")[0].childNodes[0].data
    img_h = img_size.getElementsByTagName("height")[0].childNodes[0].data
    img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
    # print("img_name:", img_name)
    # print("image_info:(w,h,c)", img_w, img_h, img_c)
    img_box = []
    for box in objects:
        cls_name = box.getElementsByTagName("name")[0].childNodes[0].data
        x1 = int(box.getElementsByTagName("xmin")[0].childNodes[0].data)
        y1 = int(box.getElementsByTagName("ymin")[0].childNodes[0].data)
        x2 = int(box.getElementsByTagName("xmax")[0].childNodes[0].data)
        y2 = int(box.getElementsByTagName("ymax")[0].childNodes[0].data)
        # print("box:(c,xmin,ymin,xmax,ymax)", cls_name, x1, y1, x2, y2)
        img_jpg_file_name = img_xml_file + '.jpg'
        img_box.append([cls_name, x1, y1, x2, y2])
    # print(img_box)

    # test_dataset_box_feature(img_jpg_file_name, img_box)
    save_file(img_xml_file, [img_w, img_h], img_box)
files = os.listdir(ANNOTATIONS_PATH)
for file in files:
    print("file name: ", file)
    file_xml = file.split(".")
    get_xml_data(ANNOTATIONS_PATH, file_xml[0])
"""
## split Images dataset
"""
from sklearn.model_selection import train_test_split
image_list = os.listdir('Dataset/FaceMask/Images')
train_list, test_list = train_test_split(image_list, test_size=0.2, random_state=7)
val_list, test_list = train_test_split(test_list, test_size=0.5, random_state=8)

print('total =',len(image_list))
print('train :',len(train_list))
print('val   :',len(val_list))
print('test  :',len(test_list))
def copy_data(file_list, img_labels_root, imgs_source, type):

    root_file = Path(DATA_ROOT + DEST_IMAGES_PATH + '/' + type)
    if not root_file.exists():
        print(f"Path {root_file} is not exit")
        os.makedirs(root_file)

    root_file = Path(DATA_ROOT + DEST_LABELS_PATH + '/' + type)
    if not root_file.exists():
        print(f"Path {root_file} is not exit")
        os.makedirs(root_file)

    for file in file_list:
        img_name = file.replace('.png', '')
        img_src_file = imgs_source + '/' + img_name + '.png'
        label_src_file = img_labels_root + '/' + img_name + '.txt'

        # print(img_sor_file)
        # print(label_sor_file)
        # im = Image.open(rf"{img_sor_file}")
        # im.show()

        # Copy image
        DICT_DIR = DATA_ROOT + DEST_IMAGES_PATH + '/' + type
        img_dict_file = DICT_DIR + '/' + img_name + '.png'

        copyfile(img_src_file, img_dict_file)

        # Copy label
        DICT_DIR = DATA_ROOT + DEST_LABELS_PATH + '/' + type
        img_dict_file = DICT_DIR + '/' + img_name + '.txt'
        copyfile(label_src_file, img_dict_file)
copy_data(train_list, LABELS_ROOT, IMAGES_ROOT, "train")
copy_data(val_list,   LABELS_ROOT, IMAGES_ROOT, "val")
copy_data(test_list,  LABELS_ROOT, IMAGES_ROOT, "test")
"""
## Create data/facemask.yaml
"""
!echo "train: Dataset/images/train\n" > data/facemask.yaml
!echo "val:   Dataset/images/val\n" >> data/facemask.yaml
!echo "nc : 3\n" >> data/facemask.yaml
!echo "names: ['With_Mask', 'Without_Mask', 'Incorrect_Mask']\n" >> data/facemask.yaml

!cat data/facemask.yaml
"""
## Train YOLOv5
<a href="../tree/yolov5/weights">link to weight folder to upload pretrained model</a>
"""
!ls
# Train with default Yolov5.weight
#!python train.py --img 320 --batch 16 --epochs 50 --data data/facemask.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt

# Train with pretrained weight >>> upload newest file to folder weights
!python train.py --img 320 --batch 16 --epochs 50 --data data/facemask.yaml --cfg models/yolov5s.yaml --weights weights/pretrained.pt

# save trained weights for detection
!cp runs/train/exp/weights/pretrained.pt weights
"""
<h2>DOWNLOAD</h2>

<a href="https://kkb-production.jupyter-proxy.kaggle.net/k/54994990/eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2IiwidHlwIjoiSldUIn0..BSWsiW8jGQTSeityClU6XQ.0zu6Bt6JSZJPCQpitQEeamc-PdEpBAyG7SYqhfSzXRaJFuqRU1HMLt9djHClz9Xygk07I4cuZxK-vAa0GjjDMQP0be1F-spgZh9WfyzQQmVPzW0eqHp8XWYRX11NiYowwjEzjCumyMKFesrvXvo_801PSCNTnbjzPrQ5aOIIjxsXf5vdBVSdRRdGtN20DOxyC2zBKpW_o583tVIYs5xxmg.JGIUJ6JicuYOERIlQcA72A/proxy/tree/yolov5/weights/best.pt"> Download BEST weight </a>
"""
"""
## Test YOLOv5
"""
"""
### detect facemask
"""
!python detect.py --source Dataset/images/test --img-size 320 --conf 0.4 --weights weights/pretrained.pt 
# display detected images
from IPython.display import Image
from glob import glob
import matplotlib.pyplot as plt
testfiles = glob('runs/detect/exp/*')

img = plt.imread(testfiles[0]) 
plt.imshow(img)    
plt.show
!python detect.py --source /kaggle/input/input-images/facemask.jpg --img-size 320 --conf 0.4 --weights weights/pretrained.pt 
Image('runs/detect/exp2/facemask.jpg')
!python detect.py --source /kaggle/input/input-images/facemask1.jpg --img-size 320 --conf 0.4 --weights weights/pretrained.pt 
Image('runs/detect/exp3/facemask1.jpg')