Spaces:
No application file
No application file
# -*- 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 |