Hugging Face's logo

Spaces:
YuXiLi
/
MaskDetection
No application file

App Files Community
MaskDetection
File size: 4,616 Bytes 
6964af9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
 
# -*- 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