Skip to content

affective_computing / face_mask_classifier

Go to a project
Toggle navigation
Toggle navigation pinning
8a4e9838 by 乔峰昇
Options

口罩二分类

0 parents
Inline Side-by-side
Showing 17 changed files with 351 additions and 0 deletions
.idea/.gitignore 0 → 100644
# Default ignored files
/shelf/
/workspace.xml
.idea/face_mask_classifier.iml 0 → 100644
<?xml version="1.0" encoding="UTF-8"?>
<module type="PYTHON_MODULE" version="4">
<component name="NewModuleRootManager">
<content url="file://$MODULE_DIR$" />
<orderEntry type="inheritedJdk" />
<orderEntry type="sourceFolder" forTests="false" />
</component>
</module>
\ No newline at end of file
.idea/inspectionProfiles/Project_Default.xml 0 → 100644
<component name="InspectionProjectProfileManager">
<profile version="1.0">
<option name="myName" value="Project Default" />
<inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
<option name="ignoredPackages">
<value>
<list size="25">
<item index="0" class="java.lang.String" itemvalue="tqdm" />
<item index="1" class="java.lang.String" itemvalue="easydict" />
<item index="2" class="java.lang.String" itemvalue="scikit_image" />
<item index="3" class="java.lang.String" itemvalue="matplotlib" />
<item index="4" class="java.lang.String" itemvalue="tensorboardX" />
<item index="5" class="java.lang.String" itemvalue="torch" />
<item index="6" class="java.lang.String" itemvalue="numpy" />
<item index="7" class="java.lang.String" itemvalue="pycocotools" />
<item index="8" class="java.lang.String" itemvalue="skimage" />
<item index="9" class="java.lang.String" itemvalue="Pillow" />
<item index="10" class="java.lang.String" itemvalue="scipy" />
<item index="11" class="java.lang.String" itemvalue="torchvision" />
<item index="12" class="java.lang.String" itemvalue="opencv_python" />
<item index="13" class="java.lang.String" itemvalue="onnxruntime" />
<item index="14" class="java.lang.String" itemvalue="onnx-simplifier" />
<item index="15" class="java.lang.String" itemvalue="onnx" />
<item index="16" class="java.lang.String" itemvalue="opencv-contrib-python" />
<item index="17" class="java.lang.String" itemvalue="numba" />
<item index="18" class="java.lang.String" itemvalue="opencv-python" />
<item index="19" class="java.lang.String" itemvalue="librosa" />
<item index="20" class="java.lang.String" itemvalue="tensorboard" />
<item index="21" class="java.lang.String" itemvalue="dill" />
<item index="22" class="java.lang.String" itemvalue="pandas" />
<item index="23" class="java.lang.String" itemvalue="scikit_learn" />
<item index="24" class="java.lang.String" itemvalue="pytorch-gradual-warmup-lr" />
</list>
</value>
</option>
</inspection_tool>
</profile>
</component>
\ No newline at end of file
.idea/inspectionProfiles/profiles_settings.xml 0 → 100644
<component name="InspectionProjectProfileManager">
<settings>
<option name="USE_PROJECT_PROFILE" value="false" />
<version value="1.0" />
</settings>
</component>
\ No newline at end of file
.idea/misc.xml 0 → 100644
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8" project-jdk-type="Python SDK" />
</project>
\ No newline at end of file
.idea/modules.xml 0 → 100644
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectModuleManager">
<modules>
<module fileurl="file://$PROJECT_DIR$/.idea/face_mask_classifier.iml" filepath="$PROJECT_DIR$/.idea/face_mask_classifier.iml" />
</modules>
</component>
</project>
\ No newline at end of file
.idea/vcs.xml 0 → 100644
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="VcsDirectoryMappings">
<mapping directory="$PROJECT_DIR$" vcs="Git" />
</component>
</project>
\ No newline at end of file
__pycache__/net.cpython-38.pyc 0 → 100644
No preview for this file type
dataset.py 0 → 100644
import os
import random
import cv2
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from utils import *
from torchvision import transforms
classes_names = ['normal', 'mask']
class FaceMaskDataset(Dataset):
def __init__(self, root_path):
self.transform = transforms.Compose([
transforms.ToTensor()
])
self.dataset = []
class_names = os.listdir(root_path)
for cls in class_names:
image_names = os.listdir(os.path.join(root_path, cls))
for image in image_names:
self.dataset.append([os.path.join(root_path, cls, image), classes_names.index(cls)])
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
lights=[0.6,0.8,1,1.2,1.4,1.6]
data = self.dataset[index]
image_path = data[0]
image_data = keep_resize_image(image_path)
image_data=cv2.convertScaleAbs(image_data,alpha=lights[random.randint(0,4)])
image_label = data[1]
return self.transform(image_data), image_label
if __name__ == '__main__':
import tqdm
d = FaceMaskDataset('image')
for i in d:
i
infer.py 0 → 100644
import os
import cv2
import torch
from PIL import Image
import numpy as np
from torchvision import transforms
from net import *
def video(net,):
cap=cv2.VideoCapture(0)
while True:
_, frame = cap.read()
image = Image.fromarray(frame)
w, h = image.size
temp = max(w, h)
mask = Image.new('RGB', (temp, temp))
if w >= h:
mask.paste(image, (0, (w - h) // 2))
else:
mask.paste(image, ((h - w) // 2, 0))
mask = mask.resize((128, 128))
mask = np.array(mask)
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)
mask_image = torch.unsqueeze(transform(mask), dim=0)
out = net(mask_image)
print(out)
out=torch.argmax(out,dim=1)
result = classes_names[int(out.item())]
cv2.putText(frame, result, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2)
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0XFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
def image_cls(net,path):
frame=cv2.imread(path)
image = Image.fromarray(frame)
w, h = image.size
temp = max(w, h)
mask = Image.new('RGB', (temp, temp))
if w >= h:
mask.paste(image, (0, (w - h) // 2))
else:
mask.paste(image, ((h - w) // 2, 0))
mask = mask.resize((128, 128))
mask = np.array(mask)
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)
mask_image = torch.unsqueeze(transform(mask), dim=0)
out = net(mask_image)
print(out)
out = torch.argmax(out, dim=1)
result = classes_names[int(out.item())]
cv2.putText(frame, result, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2)
cv2.imshow('frame', frame)
cv2.waitKey(0)
if __name__ == '__main__':
transform = transforms.Compose([
transforms.ToTensor()
])
net = FaceMaskNet()
weights_path = r'params/new_face_mobilenet_v2.pth'
classes_names = ['normal', 'mask']
if os.path.exists(weights_path):
net.load_state_dict(torch.load(weights_path, map_location='cuda:0'))
print('successfully loading weights!')
net.eval()
image_cls(net,'image/img_1.png')
net.py 0 → 100644
import torch
from torch import nn
from torchvision import models
class FaceMaskNet(nn.Module):
def __init__(self):
super(FaceMaskNet, self).__init__()
self.layer = nn.Sequential(
models.mobilenet_v2(pretrained=True)
)
self.classifier = nn.Sequential(
nn.Linear(1000, 2)
)
def forward(self, x):
return self.classifier(self.layer(x))
if __name__ == '__main__':
net = FaceMaskNet()
x = torch.randn(5, 3, 128, 128)
print(net(x).shape)
params/cls_face_mobilenet_v2.pth 0 → 100644
This file is too large to display.
test.py 0 → 100644
import os
import cv2
import torch
from net import *
from dataset import *
transform = transforms.Compose([
transforms.ToTensor()
])
data=FaceMaskDataset('/data2/face_mask')
d = DataLoader(data, batch_size=1000, shuffle=True)
with torch.no_grad():
for i,(image,label) in enumerate(d):
net=FaceMaskNet().cuda()
net.load_state_dict(torch.load('params/face_mobilenet_v2.pth'))
net.eval()
out=net(image.cuda())
out=torch.argmax(out,dim=1)
acc=torch.mean(torch.eq(label.cuda(), out).float()).item()
print(acc)
train.py 0 → 100644
import os.path
from torch import nn, optim
import torch
from dataset import *
from torch.utils.data import random_split
from net import *
import tqdm
import time
if __name__ == '__main__':
train_rate=0.8
batch_size=50
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
epochs=50
datasets = FaceMaskDataset('/data2/new_face_mask')
train_datasets, test_datasets = random_split(
datasets,
[int(len(datasets) * train_rate), len(datasets) - int(len(datasets) * train_rate)],
)
print(f'train_datasets:{len(train_datasets)} test_datasets:{len(test_datasets)}')
train_data_loader = DataLoader(train_datasets, batch_size=batch_size, shuffle=True)
test_data_loader = DataLoader(test_datasets, batch_size=batch_size, shuffle=True)
loss_fun = nn.CrossEntropyLoss()
net = FaceMaskNet().to(device)
if os.path.exists('params/new_face_mobilenet_v2.pth'):
net.load_state_dict(torch.load('params/new_face_mobilenet_v2.pth'))
print('successfully loading weights!')
opt = optim.Adam(net.parameters())
for epoch in range(1, epochs):
with tqdm.tqdm(train_data_loader) as t1:
for i, (image_data, image_label) in enumerate(train_data_loader):
net.train()
image_data, image_label = image_data.to(device), image_label.to(device)
out = net(image_data)
train_loss = loss_fun(out, image_label)
opt.zero_grad()
train_loss.backward()
opt.step()
t1.set_description(f'Epoch {epoch} train')
t1.set_postfix(train_loss=train_loss.item(),
train_acc=torch.mean(torch.eq(image_label, torch.argmax(out,dim=1)).float()).item())
time.sleep(0.1)
t1.update(1)
if (i+1) % 10 == 0:
torch.save(net.state_dict(), 'params/new_face_mobilenet_v2.pth')
print(f'epoch : {epoch} {i} successfully save weights!')
acc, temp = 0, 0
with torch.no_grad():
net.eval()
with tqdm.tqdm(test_data_loader) as t2:
for j, (image_data, image_label) in enumerate(test_data_loader):
image_data, image_label = image_data.to(device), image_label.to(device)
out = net(image_data)
test_loss = loss_fun(out, image_label)
t2.set_description(f'Epoch {epoch} test')
out = torch.argmax(out, dim=1)
t2.set_postfix(test_loss=test_loss.item(),
test_acc=torch.mean(torch.eq(image_label, out).float()).item())
time.sleep(0.1)
t2.update(1)
acc += torch.mean(torch.eq(image_label, out).float()).item()
temp += 1
print(f'epoch : {epoch} avg acc : ', acc / temp)
# acc,temp=0,0
# with torch.no_grad():
# net.eval()
# for i, (image_data, image_label) in enumerate(tqdm.tqdm(test_data_loader)):
# image_data, image_label = image_data.to(device), image_label.to(device)
# out = net(image_data)
# test_loss = loss_fun(out, image_label)
#
# out = torch.argmax(out, dim=1)
#
# acc += torch.mean(torch.eq(image_label, out).float()).item()
# temp+=1
# if i % 5 == 0:
# print(f'epoch : {epoch} {i} test_loss : ', test_loss.item())
# print(f'epoch : {epoch} {i} test acc : ',torch.mean(torch.eq(image_label, out).float()).item())
# print(f'epoch : {epoch} avg acc : ',acc/temp)
\ No newline at end of file
utils.py 0 → 100644
import cv2
import numpy as np
from PIL import Image
# 填充黑边，等比缩放
def keep_resize_image(image_path,size=(128,128)):
image=Image.open(image_path)
w,h=image.size
temp=max(w,h)
mask=Image.new('RGB',(temp,temp))
if w>=h:
mask.paste(image,(0,(w-h)//2))
else:
mask.paste(image,((h-w)//2,0))
mask=mask.resize(size)
mask=np.array(mask)
mask=cv2.cvtColor(mask,cv2.COLOR_RGB2BGR)
return mask
if __name__ == '__main__':
keep_resize_image('image/mask/img.png')
\ No newline at end of file
Styling with Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!