Skip to content

affective_computing / face_mask_classifier

Go to a project
Toggle navigation
Toggle navigation pinning
386478a4 by 乔峰昇
Options

mnn推理

1 parent 8a4e9838
Inline Side-by-side
Showing 37 changed files with 607 additions and 251 deletions
.idea/.gitignore deleted 100644 → 0
# Default ignored files
/shelf/
/workspace.xml
.idea/face_mask_classifier.iml
......@@ -2,7 +2,7 @@
<module type="PYTHON_MODULE" version="4">
<component name="NewModuleRootManager">
<content url="file://$MODULE_DIR$" />
<orderEntry type="inheritedJdk" />
<orderEntry type="jdk" jdkName="Python 3.8 (gan)" jdkType="Python SDK" />
<orderEntry type="sourceFolder" forTests="false" />
</component>
</module>
\ No newline at end of file
......
.idea/misc.xml
<?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" />
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8 (gan)" project-jdk-type="Python SDK" />
</project>
\ No newline at end of file
......
.idea/workspace.xml 0 → 100644
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ChangeListManager">
<list default="true" id="8255f694-c607-4431-a530-39f2e0df4506" name="Default Changelist" comment="">
<change beforePath="$PROJECT_DIR$/.idea/.gitignore" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/.idea/face_mask_classifier.iml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/face_mask_classifier.iml" afterDir="false" />
<change beforePath="$PROJECT_DIR$/.idea/inspectionProfiles/Project_Default.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/inspectionProfiles/Project_Default.xml" afterDir="false" />
<change beforePath="$PROJECT_DIR$/.idea/inspectionProfiles/profiles_settings.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/inspectionProfiles/profiles_settings.xml" afterDir="false" />
<change beforePath="$PROJECT_DIR$/.idea/misc.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/misc.xml" afterDir="false" />
<change beforePath="$PROJECT_DIR$/dataset.py" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/image/img.png" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/image/img_1.png" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/infer.py" beforeDir="false" afterPath="$PROJECT_DIR$/infer.py" afterDir="false" />
<change beforePath="$PROJECT_DIR$/net.py" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/params/cls_face_mobilenet_v2.pth" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/test.py" beforeDir="false" />
<change beforePath="$PROJECT_DIR$/train.py" beforeDir="false" afterPath="$PROJECT_DIR$/train.py" afterDir="false" />
<change beforePath="$PROJECT_DIR$/utils.py" beforeDir="false" />
</list>
<option name="SHOW_DIALOG" value="false" />
<option name="HIGHLIGHT_CONFLICTS" value="true" />
<option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
<option name="LAST_RESOLUTION" value="IGNORE" />
</component>
<component name="Git.Settings">
<option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
</component>
<component name="ProjectId" id="24r53vDpNxHFy2XFaKewyUFOSTL" />
<component name="ProjectViewState">
<option name="hideEmptyMiddlePackages" value="true" />
<option name="showExcludedFiles" value="false" />
<option name="showLibraryContents" value="true" />
</component>
<component name="PropertiesComponent">
<property name="RunOnceActivity.OpenProjectViewOnStart" value="true" />
<property name="RunOnceActivity.ShowReadmeOnStart" value="true" />
<property name="last_opened_file_path" value="$PROJECT_DIR$" />
<property name="settings.editor.selected.configurable" value="com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable" />
</component>
<component name="RunManager">
<configuration name="infer_mnn" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
<module name="face_mask_classifier" />
<option name="INTERPRETER_OPTIONS" value="" />
<option name="PARENT_ENVS" value="true" />
<envs>
<env name="PYTHONUNBUFFERED" value="1" />
</envs>
<option name="SDK_HOME" value="" />
<option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
<option name="IS_MODULE_SDK" value="true" />
<option name="ADD_CONTENT_ROOTS" value="true" />
<option name="ADD_SOURCE_ROOTS" value="true" />
<option name="SCRIPT_NAME" value="$PROJECT_DIR$/infer_mnn.py" />
<option name="PARAMETERS" value="" />
<option name="SHOW_COMMAND_LINE" value="false" />
<option name="EMULATE_TERMINAL" value="false" />
<option name="MODULE_MODE" value="false" />
<option name="REDIRECT_INPUT" value="false" />
<option name="INPUT_FILE" value="" />
<method v="2" />
</configuration>
<recent_temporary>
<list>
<item itemvalue="Python.infer_mnn" />
</list>
</recent_temporary>
</component>
<component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="application-level" UseSingleDictionary="true" transferred="true" />
<component name="TaskManager">
<task active="true" id="Default" summary="Default task">
<changelist id="8255f694-c607-4431-a530-39f2e0df4506" name="Default Changelist" comment="" />
<created>1644375669136</created>
<option name="number" value="Default" />
<option name="presentableId" value="Default" />
<updated>1644375669136</updated>
</task>
<servers />
</component>
</project>
\ No newline at end of file
__pycache__/net.cpython-38.pyc deleted 100644 → 0
No preview for this file type
config/config.yaml 0 → 100644
data_dir: "F:/data/face_mask" #数据集存放地址
train_rate: 0.8 #数据集划分，训练集比例
image_size: 128 #输入网络图像大小
net_type: "mobilenet_v2"
#支持模型[resnet18,resnet34,resnet50,resnet101,resnet152,resnext101_32x8d,resnext50_32x4d,wide_resnet50_2,wide_resnet101_2,
# densenet121,densenet161,densenet169,densenet201,vgg11,vgg13,vgg13_bn,vgg19,vgg19_bn,vgg16,vgg16_bn,inception_v3,
# mobilenet_v2,mobilenet_v3_small,mobilenet_v3_large,shufflenet_v2_x0_5,shufflenet_v2_x1_0,shufflenet_v2_x1_5,
# shufflenet_v2_x2_0,alexnet,googlenet,mnasnet0_5,mnasnet1_0,mnasnet1_3,mnasnet0_75,squeezenet1_0,squeezenet1_1]
# efficientnet-b0 ... efficientnet-b7
pretrained: False #是否添加预训练权重
batch_size: 50 #批次
init_lr: 0.01 #初始学习率
optimizer: 'Adam' #优化器 [SGD,ASGD,Adam,AdamW,Adamax,Adagrad,Adadelta,SparseAdam,LBFGS,Rprop,RMSprop]
class_names: [ 'mask','no_mask' ] #你的类别名称，必须和data文件夹下的类别文件名一样
epochs: 25 #训练总轮次
loss_type: "cross_entropy" # mse / l1 / smooth_l1 / cross_entropy #损失函数
model_dir: "./mobilenet_v2/weight/" #权重存放地址
log_dir: "./mobilenet_v2/logs/" # tensorboard可视化文件存放地址
\ No newline at end of file
dataset.py deleted 100644 → 0
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
import os
from PIL import Image, ImageDraw, ImageFont
import cv2
import torch
from PIL import Image
import numpy as np
from model.utils import utils
from torchvision import transforms
from net import *
from model.net.net import *
import argparse
parse = argparse.ArgumentParser('infer models')
parse.add_argument('demo', type=str, help='推理类型支持：image/video/camera')
parse.add_argument('--weights_path', type=str, default='', help='模型权重路径')
parse.add_argument('--image_path', type=str, default='', help='图片存放路径')
parse.add_argument('--video_path', type=str, default='', help='视频路径')
parse.add_argument('--camera_id', type=int, default=0, help='摄像头id')
def video(net,):
cap=cv2.VideoCapture(0)
class ModelInfer:
def __init__(self, config, weights_path):
self.config = config
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.transform = transforms.Compose([
transforms.ToTensor()
])
self.net = ClassifierNet(self.config['net_type'], len(self.config['class_names']),
False).to(self.device)
if weights_path is not None:
if os.path.exists(weights_path):
self.net.load_state_dict(torch.load(weights_path))
print('successfully loading model weights!')
else:
print('no loading model weights!')
else:
print('please input weights_path!')
exit(0)
self.net.eval()
def image_infer(self, image_path):
image = Image.open(image_path)
image_data = utils.keep_shape_resize(image, self.config['image_size'])
image_data = self.transform(image_data)
image_data = torch.unsqueeze(image_data, dim=0).to(self.device)
out = self.net(image_data)
out = torch.argmax(out)
result = self.config['class_names'][int(out)]
draw = ImageDraw.Draw(image)
font = ImageFont.truetype(r"C:\Windows\Fonts\BRITANIC.TTF", 35)
draw.text((10, 10), result, font=font, fill='red')
image.show()
def video_infer(self, video_path):
cap = cv2.VideoCapture(video_path)
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)
if _:
image_data = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image_data = Image.fromarray(image_data)
image_data = utils.keep_shape_resize(image_data, self.config['image_size'])
image_data = self.transform(image_data)
image_data = torch.unsqueeze(image_data, dim=0).to(self.device)
out = self.net(image_data)
out = torch.argmax(out)
result = self.config['class_names'][int(out)]
cv2.putText(frame, result, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), thickness=2)
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0XFF == ord('q'):
if cv2.waitKey(24) & 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()
break
image_cls(net,'image/img_1.png')
def camera_infer(self, camera_id):
cap = cv2.VideoCapture(camera_id)
while True:
_, frame = cap.read()
h, w, c = frame.shape
if _:
image_data = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image_data = Image.fromarray(image_data)
image_data = utils.keep_shape_resize(image_data, self.config['image_size'])
image_data = self.transform(image_data)
image_data = torch.unsqueeze(image_data, dim=0).to(self.device)
out = self.net(image_data)
out = torch.argmax(out)
result = self.config['class_names'][int(out)]
cv2.putText(frame, result, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), thickness=2)
cv2.imshow('frame', frame)
if cv2.waitKey(24) & 0XFF == ord('q'):
break
else:
break
if __name__ == '__main__':
args = parse.parse_args()
config = utils.load_config_util('config/config.yaml')
model = ModelInfer(config, args.weights_path)
if args.demo == 'image':
model.image_infer(args.image_path)
elif args.demo == 'video':
model.video_infer(args.video_path)
elif args.demo == 'camera':
model.camera_infer(args.camera_id)
else:
exit(0)
......
infer_mnn.py 0 → 100644
import os
import cv2
from PIL import Image, ImageFont, ImageDraw
import torch
from torchvision import transforms
import MNN
def keep_shape_resize(frame, size=128):
w, h = frame.size
temp = max(w, h)
mask = Image.new('RGB', (temp, temp), (0, 0, 0))
if w >= h:
position = (0, (w - h) // 2)
else:
position = ((h - w) // 2, 0)
mask.paste(frame, position)
mask = mask.resize((size, size))
return mask
def image_infer_mnn(mnn_model_path, image_path, class_list):
image = Image.open(image_path)
input_image = keep_shape_resize(image)
preprocess = transforms.Compose([transforms.ToTensor()])
input_data = preprocess(input_image)
interpreter = MNN.Interpreter(mnn_model_path)
session = interpreter.createSession()
input_tensor = interpreter.getSessionInput(session)
input_data = input_data.cpu().numpy().squeeze()
tmp_input = MNN.Tensor((1, 3, 128, 128), MNN.Halide_Type_Float, input_data, MNN.Tensor_DimensionType_Caffe)
input_tensor.copyFrom(tmp_input)
interpreter.runSession(session)
infer_result = interpreter.getSessionOutput(session)
output_data = infer_result.getData()
out = output_data.index(max(output_data))
draw = ImageDraw.Draw(image)
font = ImageFont.truetype(r"C:\Windows\Fonts\BRITANIC.TTF", 35)
draw.text((10, 10), class_list[int(out)], font=font, fill='red')
return image
def video_infer_mnn(mnn_model_path, video_path):
cap = cv2.VideoCapture(video_path)
interpreter = MNN.Interpreter(mnn_model_path)
session = interpreter.createSession()
input_tensor = interpreter.getSessionInput(session)
while True:
_, frame = cap.read()
if _:
image_data = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image_data = Image.fromarray(image_data)
image_data = keep_shape_resize(image_data, 128)
preprocess = transforms.Compose([transforms.ToTensor()])
input_data = preprocess(image_data)
input_data = input_data.cpu().numpy().squeeze()
tmp_input = MNN.Tensor((1, 3, 128, 128), MNN.Halide_Type_Float, input_data, MNN.Tensor_DimensionType_Caffe)
input_tensor.copyFrom(tmp_input)
interpreter.runSession(session)
infer_result = interpreter.getSessionOutput(session)
output_data = infer_result.getData()
out = output_data.index(max(output_data))
cv2.putText(frame, class_list[int(out)], (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), thickness=2)
cv2.imshow('frame', frame)
if cv2.waitKey(24) & 0XFF == ord('q'):
break
else:
break
def camera_infer_mnn(mnn_model_path, camera_id):
cap = cv2.VideoCapture(camera_id)
interpreter = MNN.Interpreter(mnn_model_path)
session = interpreter.createSession()
input_tensor = interpreter.getSessionInput(session)
while True:
_, frame = cap.read()
if _:
image_data = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image_data = Image.fromarray(image_data)
image_data = keep_shape_resize(image_data, 128)
preprocess = transforms.Compose([transforms.ToTensor()])
input_data = preprocess(image_data)
input_data = input_data.cpu().numpy().squeeze()
tmp_input = MNN.Tensor((1, 3, 128, 128), MNN.Halide_Type_Float, input_data, MNN.Tensor_DimensionType_Caffe)
input_tensor.copyFrom(tmp_input)
interpreter.runSession(session)
infer_result = interpreter.getSessionOutput(session)
output_data = infer_result.getData()
out = output_data.index(max(output_data))
cv2.putText(frame, class_list[int(out)], (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), thickness=2)
cv2.imshow('frame', frame)
if cv2.waitKey(24) & 0XFF == ord('q'):
break
else:
break
if __name__ == '__main__':
class_list = ['mask', 'no_mask']
image_path = 'test_image/mask_2997.jpg'
mnn_model_path = 'mobilenet_v2.mnn'
# image
# image=image_infer_mnn(mnn_model_path,image_path,class_list)
# image.show()
# camera
camera_infer_mnn(mnn_model_path, 0)
mobilenet_v2.mnn 0 → 100644
No preview for this file type
mobilenet_v2.onnx 0 → 100644
No preview for this file type
mobilenet_v2/weight/best.pth 0 → 100644
No preview for this file type
mobilenet_v2/weight/last.pth 0 → 100644
No preview for this file type
model/dataset/dataset.py 0 → 100644
'''
_*_coding:utf-8 _*_
@Time :2022/1/28 19:00
@Author : qiaofengsheng
@File :dataset.py
@Software :PyCharm
'''
import os
from PIL import Image
from torch.utils.data import *
from model.utils import utils
from torchvision import transforms
class ClassDataset(Dataset):
def __init__(self, data_dir, config):
self.config = config
self.transform = transforms.Compose([
transforms.RandomRotation(60),
transforms.ToTensor()
])
self.dataset = []
class_dirs = os.listdir(data_dir)
for class_dir in class_dirs:
image_names = os.listdir(os.path.join(data_dir, class_dir))
for image_name in image_names:
self.dataset.append(
[os.path.join(data_dir, class_dir, image_name),
int(config['class_names'].index(class_dir))])
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
data = self.dataset[index]
image_path, image_label = data
image = Image.open(image_path)
image = utils.keep_shape_resize(image, self.config['image_size'])
return self.transform(image), image_label
model/loss/__pycache__/loss_fun.cpython-38.pyc 0 → 100644
No preview for this file type
model/loss/loss_fun.py 0 → 100644
'''
_*_coding:utf-8 _*_
@Time :2022/1/28 19:05
@Author : qiaofengsheng
@File :loss_fun.py
@Software :PyCharm
'''
from torch import nn
class Loss:
def __init__(self, loss_type='mse'):
self.loss_fun = nn.MSELoss()
if loss_type == 'mse':
self.loss_fun = nn.MSELoss()
elif loss_type == 'l1':
self.loss_fun = nn.L1Loss()
elif loss_type == 'smooth_l1':
self.loss_fun = nn.SmoothL1Loss()
elif loss_type == 'cross_entropy':
self.loss_fun = nn.CrossEntropyLoss()
def get_loss_fun(self):
return self.loss_fun
model/net/__pycache__/net.cpython-38.pyc 0 → 100644
No preview for this file type
model/net/net.py 0 → 100644
'''
_*_coding:utf-8 _*_
@Time :2022/1/28 19:05
@Author : qiaofengsheng
@File :net.py
@Software :PyCharm
'''
import torch
from torchvision import models
from torch import nn
from efficientnet_pytorch import EfficientNet
class ClassifierNet(nn.Module):
def __init__(self, net_type='resnet18', num_classes=10, pretrained=False):
super(ClassifierNet, self).__init__()
self.layer = None
if net_type == 'resnet18': self.layer = nn.Sequential(models.resnet18(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'resnet34': self.layer = nn.Sequential(models.resnet34(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'resnet50': self.layer = nn.Sequential(models.resnet50(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'resnet101': self.layer = nn.Sequential(models.resnet101(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'resnet152': self.layer = nn.Sequential(models.resnet152(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'resnext101_32x8d': self.layer = nn.Sequential(models.resnext101_32x8d(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'resnext50_32x4d': self.layer = nn.Sequential(models.resnext50_32x4d(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'wide_resnet50_2': self.layer = nn.Sequential(models.wide_resnet50_2(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'wide_resnet101_2': self.layer = nn.Sequential(models.wide_resnet101_2(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'densenet121': self.layer = nn.Sequential(models.densenet121(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'densenet161': self.layer = nn.Sequential(models.densenet161(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'densenet169': self.layer = nn.Sequential(models.densenet169(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'densenet201': self.layer = nn.Sequential(models.densenet201(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg11': self.layer = nn.Sequential(models.vgg11(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg13': self.layer = nn.Sequential(models.vgg13(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg13_bn': self.layer = nn.Sequential(models.vgg13_bn(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg19': self.layer = nn.Sequential(models.vgg19(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg19_bn': self.layer = nn.Sequential(models.vgg19_bn(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg16': self.layer = nn.Sequential(models.vgg16(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'vgg16_bn': self.layer = nn.Sequential(models.vgg16_bn(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'inception_v3': self.layer = nn.Sequential(models.inception_v3(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mobilenet_v2': self.layer = nn.Sequential(models.mobilenet_v2(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mobilenet_v3_small': self.layer = nn.Sequential(
models.mobilenet_v3_small(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mobilenet_v3_large': self.layer = nn.Sequential(
models.mobilenet_v3_large(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'shufflenet_v2_x0_5': self.layer = nn.Sequential(
models.shufflenet_v2_x0_5(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'shufflenet_v2_x1_0': self.layer = nn.Sequential(
models.shufflenet_v2_x1_0(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'shufflenet_v2_x1_5': self.layer = nn.Sequential(
models.shufflenet_v2_x1_5(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'shufflenet_v2_x2_0': self.layer = nn.Sequential(
models.shufflenet_v2_x2_0(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'alexnet':
self.layer = nn.Sequential(models.alexnet(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'googlenet':
self.layer = nn.Sequential(models.googlenet(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mnasnet0_5':
self.layer = nn.Sequential(models.mnasnet0_5(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mnasnet1_0':
self.layer = nn.Sequential(models.mnasnet1_0(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mnasnet1_3':
self.layer = nn.Sequential(models.mnasnet1_3(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'mnasnet0_75':
self.layer = nn.Sequential(models.mnasnet0_75(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'squeezenet1_0':
self.layer = nn.Sequential(models.squeezenet1_0(pretrained=pretrained,num_classes=num_classes), )
if net_type == 'squeezenet1_1':
self.layer = nn.Sequential(models.squeezenet1_1(pretrained=pretrained,num_classes=num_classes), )
if net_type in ['efficientnet-b0', 'efficientnet-b1', 'efficientnet-b2', 'efficientnet-b3', 'efficientnet-b4',
'efficientnet-b5', 'efficientnet-b6']:
if pretrained:
self.layer = nn.Sequential(EfficientNet.from_pretrained(net_type,num_classes=num_classes))
else:
self.layer = nn.Sequential(EfficientNet.from_name(net_type,num_classes=num_classes))
def forward(self, x):
return self.layer(x)
if __name__ == '__main__':
net=ClassifierNet('mnasnet1_0',pretrained=False)
x=torch.randn(1,3,125,125)
print(net(x).shape)
model/optimizer/optim.py 0 → 100644
'''
_*_coding:utf-8 _*_
@Time :2022/1/28 19:06
@Author : qiaofengsheng
@File :optim.py
@Software :PyCharm
'''
from torch import optim
class Optimizer:
def __init__(self, net, opt_type='Adam'):
super(Optimizer, self).__init__()
self.opt = optim.Adam(net.parameters())
if opt_type == 'SGD':
self.opt = optim.SGD(net.parameters(), lr=0.01)
elif opt_type == 'ASGD':
self.opt = optim.ASGD(net.parameters())
elif opt_type == 'Adam':
self.opt = optim.Adam(net.parameters())
elif opt_type == 'AdamW':
self.opt = optim.AdamW(net.parameters())
elif opt_type == 'Adamax':
self.opt = optim.Adamax(net.parameters())
elif opt_type == 'Adagrad':
self.opt = optim.Adagrad(net.parameters())
elif opt_type == 'Adadelta':
self.opt = optim.Adadelta(net.parameters())
elif opt_type == 'SparseAdam':
self.opt = optim.SparseAdam(net.parameters())
elif opt_type == 'LBFGS':
self.opt = optim.LBFGS(net.parameters())
elif opt_type == 'Rprop':
self.opt = optim.Rprop(net.parameters())
elif opt_type == 'RMSprop':
self.opt = optim.RMSprop(net.parameters())
def get_optimizer(self):
return self.opt
model/utils/__pycache__/utils.cpython-38.pyc 0 → 100644
No preview for this file type
model/utils/utils.py 0 → 100644
'''
_*_coding:utf-8 _*_
@Time :2022/1/28 19:58
@Author : qiaofengsheng
@File :utils.py
@Software :PyCharm
'''
import torch
import yaml
from PIL import Image
from torch.nn.functional import one_hot
def load_config_util(config_path):
config_file = open(config_path, 'r', encoding='utf-8')
config_data = yaml.load(config_file)
return config_data
def keep_shape_resize(frame, size=128):
w, h = frame.size
temp = max(w, h)
mask = Image.new('RGB', (temp, temp), (0, 0, 0))
if w >= h:
position = (0, (w - h) // 2)
else:
position = ((h - w) // 2, 0)
mask.paste(frame, position)
mask = mask.resize((size, size))
return mask
def label_one_hot(label):
return one_hot(torch.tensor(label))
net.py deleted 100644 → 0
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 deleted 100644 → 0
This file is too large to display.
test.py deleted 100644 → 0
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
import os.path
from torch import nn, optim
import os.path
import time
import torch
from dataset import *
from torch.utils.data import random_split
from net import *
import tqdm
import time
from torch.utils.tensorboard import SummaryWriter
from model.net.net import ClassifierNet
from model.loss.loss_fun import *
from model.optimizer.optim import *
from model.dataset.dataset import *
import argparse
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
parse = argparse.ArgumentParser(description='train_demo of argparse')
parse.add_argument('--weights_path', default=None)
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())
class Train:
def __init__(self, config):
self.config = config
if not os.path.exists(config['model_dir']):
os.makedirs(config['model_dir'])
self.summary_writer = SummaryWriter(config['log_dir'])
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net = ClassifierNet(self.config['net_type'], len(self.config['class_names']),
self.config['pretrained']).to(self.device)
self.loss_fun = Loss(self.config['loss_type']).get_loss_fun()
self.optimizer = Optimizer(self.net, self.config['optimizer']).get_optimizer()
self.dataset = ClassDataset(self.config['data_dir'], config)
self.train_dataset, self.test_dataset = random_split(self.dataset,
[int(len(self.dataset) * config['train_rate']),
len(self.dataset) - int(
len(self.dataset) * config['train_rate'])]
)
self.train_data_loader = DataLoader(self.train_dataset, batch_size=self.config['batch_size'], shuffle=True)
self.test_data_loader = DataLoader(self.test_dataset, batch_size=self.config['batch_size'], shuffle=True)
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())
def train(self, weights_path):
print(f'device:{self.device} 训练集:{len(self.train_dataset)} 测试集:{len(self.test_dataset)}')
if weights_path is not None:
if os.path.exists(weights_path):
self.net.load_state_dict(torch.load(weights_path))
print('successfully loading model weights!')
else:
print('no loading model weights')
temp_acc = 0
for epoch in range(1, self.config['epochs'] + 1):
self.net.train()
with tqdm.tqdm(self.train_data_loader) as t1:
for i, (image_data, image_label) in enumerate(self.train_data_loader):
image_data, image_label = image_data.to(self.device), image_label.to(self.device)
out = self.net(image_data)
if self.config['loss_type'] == 'cross_entropy':
train_loss = self.loss_fun(out, image_label)
else:
train_loss = self.loss_fun(out, utils.label_one_hot(image_label).type(torch.FloatTensor).to(
self.device))
t1.set_description(f'Train-Epoch {epoch} 轮 {i} 批次 : ')
t1.set_postfix(train_loss=train_loss.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!')
self.optimizer.zero_grad()
train_loss.backward()
self.optimizer.step()
if i % 10 == 0:
torch.save(self.net.state_dict(), os.path.join(self.config['model_dir'], 'last.pth'))
self.summary_writer.add_scalar('train_loss', train_loss.item(), epoch)
self.net.eval()
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')
with tqdm.tqdm(self.test_data_loader) as t2:
for j, (image_data, image_label) in enumerate(self.test_data_loader):
image_data, image_label = image_data.to(self.device), image_label.to(self.device)
out = self.net(image_data)
if self.config['loss_type'] == 'cross_entropy':
test_loss = self.loss_fun(out, image_label)
else:
test_loss = self.loss_fun(out, utils.label_one_hot(image_label).type(torch.FloatTensor).to(
self.device))
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())
test_acc = torch.mean(torch.eq(out, image_label).float()).item()
acc += test_acc
temp += 1
t2.set_description(f'Test-Epoch {epoch} 轮 {j} 批次 : ')
t2.set_postfix(test_loss=test_loss.item(), test_acc=test_acc)
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)
print(f'Test-Epoch {epoch} 轮准确率为 : {acc / temp}')
if (acc / temp) > temp_acc:
temp_acc = acc / temp
torch.save(self.net.state_dict(), os.path.join(self.config['model_dir'], 'best.pth'))
else:
torch.save(self.net.state_dict(), os.path.join(self.config['model_dir'], 'last.pth'))
self.summary_writer.add_scalar('test_loss', test_loss.item(), epoch)
self.summary_writer.add_scalar('test_acc', acc / temp, epoch)
# 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
if __name__ == '__main__':
args = parse.parse_args()
config = utils.load_config_util('config/config.yaml')
Train(config).train(args.weights_path)
......
utils.py deleted 100644 → 0
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!