# -*- coding: utf-8 -*-
"""DL_finalProject.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/10aj1WXX-L9ZzCbSXCLN1D3X6Pwbf1OCr
"""

# Commented out IPython magic to ensure Python compatibility.
# 連結google drive
from google.colab import drive
drive.mount('/content/drive/')
# %cd /content/drive/MyDrive/AI/DL_finalProject

!pip install -q kaggle

# 上傳kaggle.json
from google.colab import files
uploaded = files.upload()

!mkdir ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json

import kaggle
# download API
!kaggle datasets download -d andrewmvd/face-mask-detection

# 將檔案解壓縮並放至指定資料夾
!unzip face-mask-detection.zip

# 載入與創建資料夾
from pathlib import Path
images_dir= Path('images')
annotations_dir= Path('annotations')

labels_dir = Path('labels')
if not labels_dir.exists():
  !mkdir -p labels

import xml.etree.ElementTree as ET
import os

# 將bbox轉換成yolov7可訓練的格式
def convert_bbox(size, box):
  dw = 1. / size[0]
  dh = 1. / size[1]
  x = (box[0] + box[1]) / 2.0
  y = (box[2] + box[3]) / 2.0
  w = box[1] - box[0]
  h = box[3] - box[2]
  return (x * dw, y * dh, w * dw, h * dh)

classes = {'without_mask': 0 ,'mask_weared_incorrect': 1 ,'with_mask': 2}

# 遍歷annotations文件夾中的所有.xml文件
for xml_file in os.listdir(annotations_dir):
  if not xml_file.endswith('.xml'):
    continue
  tree = ET.parse(os.path.join(annotations_dir, xml_file))
  root = tree.getroot()

  # 擷取圖片寬度及高度資訊
  size = root.find('size')
  w = int(size.find('width').text)
  h = int(size.find('height').text)

  # 創建對應的txt文件
  txt_file = xml_file.replace('.xml', '.txt')
  with open(os.path.join(labels_dir, txt_file), 'w') as f:
    for obj in root.iter('object'):
        # 讀取類別名稱
        cls = obj.find('name').text
        # 讀取bbox資訊
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
                float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert_bbox((w, h), b)
        f.write(f"{classes[cls]} {bb[0]} {bb[1]} {bb[2]} {bb[3]}\n")

import os, shutil, random

# 設置比例
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

# 準備文件夾
base_dir = 'data'
if not os.path.exists(base_dir):
  os.makedirs(base_dir, exist_ok=True)

for subset in ['train', 'val', 'test']:
  for folder in ['images', 'labels']:
    os.makedirs(f"{base_dir}/{subset}/{folder}", exist_ok=True)

# 抓取圖片檔名（maksssksksss___.png）並打亂順序
image_files = [f for f in os.listdir(images_dir)]
random.seed(603)
random.shuffle(image_files)

# 計算集合大小
total = len(image_files)
train_size = int(total * train_ratio)
val_size = int(total * val_ratio)
print('train:', train_size, 'val:', val_size, 'test:', total-train_size-val_size)

# 分割數據集
train_files = image_files[:train_size]
val_files = image_files[train_size:train_size + val_size]
test_files = image_files[train_size + val_size:]

# 移動文件
for subset, files in [('train', train_files), ('val', val_files), ('test', test_files)]:
  for file in files:
    # images
    src_img = os.path.join(images_dir, file)
    dst_img = os.path.join(f"{base_dir}/{subset}/images", file)
    shutil.copy(src_img, dst_img)

    # labels
    label_file = file.rsplit('.', 1)[0] + '.txt'
    src_label = os.path.join(labels_dir, label_file)
    dst_label = os.path.join(f"{base_dir}/{subset}/labels", label_file)
    shutil.copy(src_label, dst_label)

# 寫入data.yaml
yaml= 'data/data.yaml'

with open(yaml, 'w') as f:
    f.write(f"""
train: ../data/train
val: ../data/val
test: ../data/test

nc: 3
names: ['without_mask', 'mask_weared_incorrect', 'with_mask']
""")

# Commented out IPython magic to ensure Python compatibility.
# 下載YOLOv7
!git clone https://github.com/WongKinYiu/yolov7
# %cd yolov7/
!pip install -r requirements.txt

# Commented out IPython magic to ensure Python compatibility.
# %cd yolov7
# 模型訓練
!python train.py --img 640 \
          --batch 16 \
          --epochs 80 \
          --data ../data/data.yaml \
          --cfg cfg/training/yolov7.yaml \
          --weights yolov7.pt \
          --name results \
          --device 0 \

# Commented out IPython magic to ensure Python compatibility.
# 模型測試
# %cd yolov7
!python test.py --data ../data/data.yaml\
         --weights runs/train/results/weights/best.pt\
         --batch-size 4\
         --conf 0.2\
         --img-size 640\
         --task test