def creat_mtcnn_detector(prefix, epoch, batch_size, test_mode, thresh,
min_face_size, ctx):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
PNet = FcnDetector(P_Net20("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
if test_mode in ["onet", "hardrnet", "hardonet"]:
args, auxs = load_param(prefix[1], epoch[1], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
if test_mode == "hardonet":
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test", False), 48, batch_size[2], ctx, args,
auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=4,
threshold=thresh,
slide_window=False)
return mtcnn_detector
def mtcnn_model(prefix, epoch, batch_size, ctx, thresh, min_face, stride,
slide_window):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face,
stride=stride,
threshold=thresh,
slide_window=slide_window)
return mtcnn_detector
def test_net(gpuId):
prefix = [
os.path.join(mtnnDir, 'pnet'),
os.path.join(mtnnDir, 'rnet'),
os.path.join(mtnnDir, 'onet')
]
epoch = [16, 16, 16]
batch_size = [2048, 256, 16]
ctx = mx.gpu(gpuId)
thresh = [0.5, 0.5, 0.7]
min_face_size = 40
stride = 2
args_p, auxs_p = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
PNet = FcnDetector(P_Net("test"), ctx, args_p, auxs_p)
# load rnet model
args_r, auxs_r = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args_r, auxs_r)
# load onet model
args_o, auxs_o = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args_o, auxs_o)
return MtcnnDetector(
detectors=[PNet, RNet, ONet],
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=False)
def test_net(root_path,
dataset_path,
image_set,
prefix,
epoch,
batch_size,
test_mode="rnet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
shuffle=False,
vis=False):
detectors = [None, None, None]
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
# load pnet model
if slide_window:
PNet = Detector(P_Net, 12, batch_size[0], model_path[0])
else:
PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
RNet = Detector(R_Net, 24, batch_size[1], model_path[1])
detectors[1] = RNet
# load onet model
if test_mode == "onet":
ONet = Detector(O_Net, 48, batch_size[2], model_path[2])
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
imdb = IMDB("wider", image_set, root_path, dataset_path, 'test')
gt_imdb = imdb.gt_imdb()
test_data = TestLoader(gt_imdb)
detections = mtcnn_detector.detect_face(imdb, test_data, vis=vis)
if test_mode == "pnet":
net = "rnet"
elif test_mode == "rnet":
net = "onet"
save_path = os.path.join(data_dir, net)
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections.pkl")
with open(save_file, 'wb') as f:
cPickle.dump(detections, f, cPickle.HIGHEST_PROTOCOL)
save_hard_example(net)
def test_net(root_path,
dataset_path,
image_set,
prefix,
epoch,
batch_size,
ctx,
test_mode="onet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
shuffle=False,
vis=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
args, auxs = load_param(prefix[1], epoch[1], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
if test_mode == "onet":
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
imdb = IMDB("fddb", image_set, root_path, dataset_path, 'test')
gt_imdb = imdb.gt_imdb()
test_data = TestLoader(gt_imdb)
print('1')
# 得到的detections格式为:[[第一张图的所有预测框],[第二张图片的所有预测框],[第三张图片的所有预测框],...,[最后一张图片的所有预测框]],每个预测框为[x1,y1,x2,y2,score],score为该预测框为人脸的分数
detections = mtcnn_detector.detect_face(imdb, test_data, vis=vis)
save_path = "/Users/qiuxiaocong/Downloads/mtcnn1"
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections_onet_0009_givenPRnet.pkl")
with open(save_file, 'wb') as f:
pickle.dump(detections, f, pickle.HIGHEST_PROTOCOL)
print('detections saved done!')
def test_net(prefix,
epoch,
batch_size,
ctx,
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=4,
slide_window=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net20("test"), 20, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net20("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[1], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
img = cv2.imread('test01.jpg')
t1 = time.time()
boxes, boxes_c = mtcnn_detector.detect_pnet20(img)
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
print 'time: ', time.time() - t1
if boxes_c is not None:
draw = img.copy()
font = cv2.FONT_HERSHEY_SIMPLEX
for b in boxes_c:
cv2.rectangle(draw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])),
(0, 255, 255), 1)
cv2.putText(draw, '%.3f' % b[4], (int(b[0]), int(b[1])), font, 0.4,
(255, 255, 255), 1)
cv2.imshow("detection result", draw)
cv2.waitKey(0)
def test_net(root_path,
dataset_path,
prefix,
epoch,
batch_size,
ctx,
test_mode="onet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
shuffle=False,
vis=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=False, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
args, auxs = load_param(prefix[1], epoch[0], convert=False, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
if test_mode == "onet":
args, auxs = load_param(prefix[2], epoch[2], convert=False, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
for i in range(1, 11):
image_set = "fold-" + str(i).zfill(2)
imdb = IMDB("fddb", image_set, root_path, dataset_path, 'test')
gt_imdb = imdb.gt_imdb()
test_data = TestLoader(gt_imdb)
all_boxes = mtcnn_detector.detect_face(imdb, test_data, vis=vis)
imdb.write_results(all_boxes)
def test_net(root_path, dataset_path, image_set, prefix, epoch,
batch_size, ctx, test_mode="rnet",
thresh=[0.6, 0.6, 0.7], min_face_size=24,
stride=2, slide_window=False, shuffle=False, vis=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
args, auxs = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
if test_mode == "onet":
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors, ctx=ctx, min_face_size=min_face_size,
stride=stride, threshold=thresh, slide_window=slide_window)
imdb = IMDB("wider", image_set, root_path, dataset_path, 'test')
gt_imdb = imdb.gt_imdb()
test_data = TestLoader(gt_imdb)
detections = mtcnn_detector.detect_face(imdb, test_data, vis=vis)
if test_mode == "pnet":
net = "rnet"
elif test_mode == "rnet":
net = "onet"
save_path = "./prepare_data/%s"%net
if not os.path.exists(save_path):
os.mkdir(save_path)
save_file = os.path.join(save_path, "detections.pkl")
with open(save_file, 'wb') as f:
cPickle.dump(detections, f, cPickle.HIGHEST_PROTOCOL)
save_hard_example(net)
def test_net(root_path,
dataset_path,
prefix,
epoch,
batch_size,
test_mode="onet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
shuffle=False,
vis=False):
detectors = [None, None, None]
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
# load pnet model
if slide_window:
PNet = Detector(P_Net, 12, batch_size[0], model_path[0])
else:
PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
RNet = Detector(R_Net, 24, batch_size[1], model_path[1])
detectors[1] = RNet
# load onet model
if test_mode == "onet":
ONet = Detector(O_Net, 48, batch_size[2], model_path[2])
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
for i in range(1, 11):
image_set = "fold-" + str(i).zfill(2)
imdb = IMDB("fddb", image_set, root_path, dataset_path, 'test')
gt_imdb = imdb.gt_imdb()
test_data = TestLoader(gt_imdb)
all_boxes = mtcnn_detector.detect_face(imdb, test_data, vis=vis)
imdb.write_results(all_boxes)
def main():
image_path = sys.argv[1]
casc_path = sys.argv[2]
image = Image(image_path)
image_loaded = image.get_image()
preprocessor = Preprocessor(image_loaded)
image_processed = preprocessor.preprocess()
detector = Detector(casc_path)
faces = detector.detect(image_processed)
print "Found {0} faces!".format(len(faces))
print "Faces: ", faces
visualizer = Visualizer(image.get_image())
visualizer.viz(faces)
cv2.waitKey(0)
def openImageHandler(self, event):
filename = askopenfilename(title="Choose an image")
self.sourceImage.canvas.delete(ALL)
img = Image.open(filename)
basewidth = 400
wpercent = (basewidth/float(img.size[0]))
hsize = int((float(img.size[1])*float(wpercent)))
img = img.resize((basewidth,hsize), Image.ANTIALIAS)
self.photo = ImageTk.PhotoImage(img)
self.sourceImage.canvas.create_image(200, 200, image=self.photo)
self.detector = Detector(filename)
img = Image.fromarray(self.detector.detect_image())
wpercent = (basewidth/float(img.size[0]))
hsize = int((float(img.size[1])*float(wpercent)))
img = img.resize((basewidth,hsize), Image.ANTIALIAS)
self.photo_result = ImageTk.PhotoImage(img)
self.detectionImage.canvas.create_image(200, 200, image=self.photo_result)
self.initDetector()
def load_model(model_path):
"""
根据指定的模型路径加载mtcnn和DenseNet的模型
:param model_path:
:return:
"""
batch_size = [1, 1, 1]
model = DenseNet(24,
40,
3,
0.8,
'DenseNet-BC',
reduction=0.5,
bc_mode=True)
detectors = [None, None, None]
# 加载用于分类的DenseNet模型
model.load_model()
# 计算出PNet、RNet、ONet模型的路径
prefix = ['pnet/pnet', 'rnet/rnet', 'onet/onet']
prefix = [os.path.join(model_path, x) for x in prefix]
epoch = [7, 7, 7]
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
# 加载PNet
p_net = FcnDetector(P_Net, model_path[0])
detectors[0] = p_net
# 加载RNet
r_net = Detector(R_Net, 24, batch_size[1], model_path[1])
detectors[1] = r_net
# 加载ONet
o_net = Detector(O_Net, 48, batch_size[2], model_path[2])
detectors[2] = o_net
# 建立MTCNN检测器
mtcnn_detector = MtcnnDetector(detectors=detectors,
min_face_size=24,
stride=2,
threshold=(0.6, 0.6, 0.7),
slide_window=False)
return model, mtcnn_detector
def test(prefix,
epoch,
batch_size,
test_mode="onet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
shuffle=False,
vis=False):
#img = cv2.imread("./1111.jpg")
detectors = [None, None, None]
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
print(model_path)
# load pnet model
if slide_window:
PNet = Detector(P_Net, 12, batch_size[0], model_path[0])
else:
PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
RNet = Detector(R_Net, 24, batch_size[1], model_path[1])
detectors[1] = RNet
# load onet model
if test_mode == "onet":
ONet = Detector(O_Net, 48, batch_size[2], model_path[2])
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
for name in os.listdir("C:\\Users\\JINNIU\\Desktop\\liuzhen\\qinghua"):
img = cv2.imread(
os.path.join("C:\\Users\\JINNIU\\Desktop\\liuzhen\\qinghua", name))
boxes, boxes_c = mtcnn_detector.detect(img)
visssss(img, boxes_c, name, thresh=0.998)
def run_detector(args, zmq_endpoint, log_file):
camera_img_width = 1920
camera_img_height = 1080
try:
detector = Detector(
camera_ip=args.camera_ip,
intersection=args.intersection,
train_detection_model_weights=args.train_model_weights,
signal_detection_model_weights=args.signal_model_weights,
camera_img_width=camera_img_width,
camera_img_height=camera_img_height,
log_file=log_file,
sleep_length=args.sleep_length,
zmq_endpoint=zmq_endpoint,
contrast_pp_alpha=args.contrast_alpha,
contrast_pp_threshold=args.contrast_threshold)
detector.run()
except:
traceback.print_exc(file=sys.stdout)
sys.stdout.flush()
def loadModel(self):
self.model = face_embedding.FaceModel(self.args)
detectors = [None, None, None]
ctx = mx.gpu(0)
prefix = ['mtcnnmodel/pnet', 'mtcnnmodel/rnet', 'mtcnnmodel/onet']
epoch = [16, 16, 16]
batch_size = [2048, 256, 16]
thresh = [0.6, 0.6, 0.7]
min_face_size = 24
stride = 2
slide_window = False
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
self.mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
#print (self.model)
self.id_dataset, self.idnums = self.get_id_data(self.args.id_dir)
cap = cv2.VideoCapture(0 + cv2.CAP_DSHOW) # 参数为0时调用本地摄像头；url连接调取网络摄像头；文件地址获取本地视频
cap.set(3, 1920) # 设置分辨率
cap.set(4, 1080)
cap.set(cv2.CAP_PROP_FPS, 30)
ret, frame = cap.read()
while (True):
ret, frame = cap.read()
frame = frame[:, ::-1]
frame = frame[:, 440: -440]
image = cv2.resize(frame, (size, size))
bboxes = det.predict(image.copy(), size, (0.5, 0.5))
for cid, bbox in bboxes[0].items():
cls = "mask" if cid == 1 else "face"
for b in bbox:
prob = b[-1]
b = b[:4].astype(int)
cv2.rectangle(image, (b[0], b[1]), (b[2], b[3]), colors[cid].tolist(), 1, cv2.LINE_AA)
cv2.putText(image, "{}:{}".format(cls, int(prob * 100)), (b[0], b[1]), cv2.FONT_ITALIC, 1,
colors[cid].tolist(), 2)
cv2.imshow("image", image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if __name__ == '__main__':
det = Detector(classes_info, model_info, "cuda")
det.load_model("checkpoints/2021-03-08 00.11.56/epoch=331_4.7689.pth")
# to_image(det)
to_video(det)
def test_net(prefix,
epoch,
batch_size,
ctx,
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window: # 使用滑动窗口(MTCNN的P_NET不使用了滑动窗口,而是全卷积网络)
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[1], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector1(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
# img = cv2.imread('test01.jpg') # 读取图片
# img = cv2.imread('zhang.jpeg') # 读取图片
# img = cv2.imread('curry.jpg') # 读取图片
# img = cv2.imread('physics.jpg') # 读取图片
# img = cv2.imread('000007.jpg') # 读取图片
# img = cv2.imread('test01.jpg') # 读取图片
# img = cv2.imread('NBA98.jpg')
# img = cv2.imread('download.jpg')
# img = cv2.imread('/Users/qiuxiaocong/Downloads/WIDER_train/images/7--Cheering/7_Cheering_Cheering_7_16.jpg')
# img = cv2.imread('/Users/qiuxiaocong/Downloads/WIDER_train/images/11--Meeting/11_Meeting_Meeting_11_Meeting_Meeting_11_77.jpg')
# img = cv2.imread('/Users/qiuxiaocong/Downloads/3Dfacedeblurring/dataset_test/falling1/input/00136.png')
img = cv2.imread('/Users/qiuxiaocong/Downloads/facetrack_python/error.jpg')
boxes, boxes_c = mtcnn_detector.detect_pnet(img)
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
# print(boxes_c) # x1 y1 x2 y2
original_detect = []
crop_list = []
detect_len_list = []
nd_array = []
score_list = []
if boxes_c is not None:
draw = img.copy()
font = cv2.FONT_HERSHEY_SIMPLEX # Python 一种字体
idx = 0
for b in boxes_c: # nms和bbr之后的结果
# 在draw上绘制矩形框(左上角坐标+右下角坐标)
b_new0 = np.array(b[0:4]) # 添加检测框
original_detect.append(b_new0)
b_new = convert_to_square(b_new0) # 添加送入到landmark net的48*48大小的框
crop_list.append(b_new)
score_list.append(b[4])
# cv2.rectangle(draw, (int(b_new[0]), int(b_new[1])), (int(b_new[2]), int(b_new[3])),
# (0, 255, 255), 1)
# # 在draw上添加文字
# cv2.putText(draw, '%.3f'%b[4], (int(b[0]), int(b[1])), font, 0.4, (255, 255, 255), 1)
# cv2.imshow("detection result", draw)
img_draw = img[int(b_new[1]):int(b_new[3]),
int(b_new[0]):int(b_new[2])]
detect_len = min(img_draw.shape[0], img_draw.shape[1])
# print(img_draw.shape[0], img_draw.shape[1])
if detect_len != 0:
detect_len_list.append(detect_len)
img_resized = cv2.resize(img_draw, (48, 48))
# cv2.imshow("detection result", draw)
# print('img_resized type is :{}'.format(type(img_resized)))
nd_array.append(img_resized)
# cv2.imwrite("detection_result{}.jpg".format(idx), img_resized)
# cv2.waitKey(0)
idx = idx + 1
return crop_list, detect_len_list, original_detect, idx, img, nd_array
class FuncButtons:
def __init__(self, master, side=None, sourceImage=None, detectionImage=None, textBoxDetectionResult=None, textMatchedFacts=None, textHitRules = None):
frame = Frame(master)
if side:
frame.pack(side=side)
else:
frame.pack()
self.sourceImage = sourceImage
self.detectionImage = detectionImage
self.textBoxDetectionResult=textBoxDetectionResult
self.textMatchedFacts=textMatchedFacts
self.textHitRules = textHitRules
self.openImageButton = CustomButton(frame, text="Open Image", handler=self.openImageHandler)
self.openRuleEditorButton = CustomButton(frame, text="Open Rule Editor", handler=self.openRuleEditorHandler)
self.showRulesButton = CustomButton(frame, text="Show Rules", handler=self.showRulesHandler)
self.showFactsButton = CustomButton(frame, text="Show Facts", handler=self.showFactsHandler)
def openImageHandler(self, event):
filename = askopenfilename(title="Choose an image")
self.sourceImage.canvas.delete(ALL)
img = Image.open(filename)
basewidth = 400
wpercent = (basewidth/float(img.size[0]))
hsize = int((float(img.size[1])*float(wpercent)))
img = img.resize((basewidth,hsize), Image.ANTIALIAS)
self.photo = ImageTk.PhotoImage(img)
self.sourceImage.canvas.create_image(200, 200, image=self.photo)
self.detector = Detector(filename)
img = Image.fromarray(self.detector.detect_image())
wpercent = (basewidth/float(img.size[0]))
hsize = int((float(img.size[1])*float(wpercent)))
img = img.resize((basewidth,hsize), Image.ANTIALIAS)
self.photo_result = ImageTk.PhotoImage(img)
self.detectionImage.canvas.create_image(200, 200, image=self.photo_result)
self.initDetector()
def initDetector(self):
fact_templates = [
'sisi', 'jumlahSisiSama', 'sudutTumpul', 'sudutLancip', 'sudutSiku',
'jumlahPasanganSejajar', 'sepasangSudutSama', 'maksXSiku', 'minXSiku', 'jumlahSudutSama'
]
fact_results = []
delta = 5
fact_results.append(self.detector.get_edges_length())
# Sisi
lengths = self.detector.get_all_length()
# jumlah sisi sama
count = 0
for i in range(len(lengths)):
temp_count = 0
for j in range (len(lengths)):
if (lengths[i] >= lengths[j]-delta and lengths[i] <= lengths[j] + delta):
temp_count += 1
if temp_count > count:
count = temp_count
fact_results.append(count)
degrees = self.detector.get_all_degrees()
count_tumpul = 0
count_siku = 0
count_lancip = 0
for degree in degrees:
if (degree > 90-delta and degree < 90+delta):
count_siku += 1
elif (degree <= 90-delta):
count_lancip += 1
else:
count_tumpul += 1
# Sudut tumpul
fact_results.append(count_tumpul)
# Sudut lancip
fact_results.append(count_lancip)
# Sudut Siku
fact_results.append(count_siku)
# jumlahPasanganSejajar
slopes = self.detector.get_all_slopes()
count = 0
forbidden_index = []
for i in range (len(slopes)):
for j in range(i+1, len(slopes)):
if (slopes[i] >= slopes[j] - delta and slopes[i] <= slopes[j] + delta and i not in forbidden_index and j not in forbidden_index):
count += 1
forbidden_index.append(i)
forbidden_index.append(j)
fact_results.append(count)
# Sepasang sudut sama
count = 0
forbidden_index = []
for i in range (len(degrees)):
for j in range(i+1, len(degrees)):
if (degrees[i] >= degrees[j] - delta and degrees[i] <= degrees[j] + delta and i not in forbidden_index and j not in forbidden_index):
count += 1
forbidden_index.append(i)
forbidden_index.append(j)
fact_results.append(count)
# maksXsiku
edges = self.detector.get_all_edges()
max_index = 0
min_index = 0
for i in range(len(edges)):
if (edges[i][0] > edges[max_index][0]):
max_index = i
if (edges[i][0] < edges[min_index][0]):
min_index = i
if (degrees[max_index] > 90-delta and degrees[max_index] < 90+delta):
fact_results.append('true')
else:
fact_results.append('false')
#minXsiku
if (degrees[min_index] > 90-delta and degrees[min_index] < 90+delta):
fact_results.append('true')
else:
fact_results.append('false')
# jumlahsudutsama
count = 0
for i in range(len(degrees)):
temp_count = 0
for j in range (len(degrees)):
if (degrees[i] == degrees[j]):
temp_count += 1
if temp_count > count:
count = temp_count
fact_results.append(count)
core = Core('rules.clp')
for i in range(10):
fact = '('+str(fact_templates[i])+' '+str(fact_results[i])+')'
core.assert_fact(fact)
# Run and get hit rules
hit_rules = core.get_hit_rules()
# Get results
object_result = core.get_results()
# Get all matched facts
facts = core.get_matched_facts()
# print(hit_rules)
# print(str(object_result))
# print(facts)
# Write to text box
self.textHitRules.clearTextBox()
for hit_rule in hit_rules:
self.textHitRules.insertTextLine(hit_rule)
self.textBoxDetectionResult.clearTextBox()
self.textBoxDetectionResult.insertTextLine(object_result)
self.textMatchedFacts.clearTextBox()
for fact in facts:
self.textMatchedFacts.insertTextLine(fact)
def openRuleEditorHandler(self, event):
webbrowser.open("rules.clp")
def showRulesHandler(self, event):
self.textHitRules.showTextBox();
def showFactsHandler(self, event):
self.textMatchedFacts.showTextBox();
def test_net(prefix,
epoch,
batch_size,
ctx,
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
video_capture = cv2.VideoCapture(0)
boxes = []
boxes_c = []
while True:
#img = cv2.imread('/home/zzg/Opensource/mtcnn-master/data/custom/02.jpg')
_, img = video_capture.read()
t1 = time.time()
boxes, boxes_c = mtcnn_detector.detect_pnet(img)
if boxes_c is None:
continue
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
if boxes_c is None:
continue
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
print 'time: ', time.time() - t1
if boxes_c is not None:
draw = img.copy()
font = cv2.FONT_HERSHEY_SIMPLEX
for b in boxes_c:
cv2.rectangle(draw, (int(b[0]), int(b[1])),
(int(b[2]), int(b[3])), (0, 255, 255), 1)
#cv2.putText(draw, '%.3f'%b[4], (int(b[0]), int(b[1])), font, 0.4, (255, 255, 255), 1)
cv2.imshow("detection result", draw)
#cv2.imwrite("o12.jpg",draw)
cv2.waitKey(10)
def detectOneImg(prefix,
epoch,
batch_size,
model,
imgPath,
test_mode="onet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
shuffle=False,
vis=False):
detectors = [None, None, None]
#load densenet for classfication
model.load_model()
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
print(model_path)
# load pnet model
if slide_window:
PNet = Detector(P_Net, 12, batch_size[0], model_path[0])
else:
PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet
# load rnet model
if test_mode in ["rnet", "onet"]:
RNet = Detector(R_Net, 24, batch_size[1], model_path[1])
detectors[1] = RNet
# load onet model
if test_mode == "onet":
ONet = Detector(O_Net, 48, batch_size[2], model_path[2])
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
images_res = []
img = cv2.imread(os.path.join(imgPath))
boxes, boxes_c = mtcnn_detector.detect(img)
#print(boxes_c.shape)
#visssss(img, boxes_c, 'plain13134.jpg', thresh=0.998)
for i in range(boxes_c.shape[0]):
bbox = boxes_c[i, :4].astype('int32')
if bbox[1] < 0:
bbox[1] = 0
if bbox[0] < 0:
bbox[0] = 0
if bbox[2] > 2048:
bbox[2] = 2048
if bbox[3] > 2048:
bbox[3] = 2048
crop = img[bbox[1]:bbox[3], bbox[0]:bbox[2], :]
#print(boxes_c[i, :4])
#print('crop:',crop.shape)
crop = cv2.resize(crop, (48, 48))
#cv2.imwrite("C:\\Users\\JINNIU\\Desktop\\liuzhen\\qinghua\\temp\\"+str(i)+'.jpg',crop)
images_res.append(crop)
images_res = np.array(images_res).astype(np.float32)
images_res = normalize_images(images_res)
pred = model.test(images_res)
bg_box = np.where(
pred == 45
) #if the class is 45, the image is the background and not the traffic sign. we omit these in the next step
#print(pred)
#print(len(boxes_cc))
for ii in bg_box[0]:
boxes_c[ii, :] = 0
#print(boxes_c)
#del boxes_c[np.where(pred==45),:]
visssss(img, boxes_c, imgPath, pred, thresh=0.998)
def test_net(prefix=['model/pnet', 'model/rnet', 'model/onet'],
epoch=[16, 16, 16],
batch_size=[2048, 256, 16],
ctx=mx.cpu(0),
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
camera_path='0'):
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
# load rnet model
args, auxs = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
mtcnn_detector = MtcnnDetector(detectors=[PNet, RNet, ONet],
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=False)
try:
capture = cv2.VideoCapture(int(camera_path))
except ValueError as e:
capture = cv2.VideoCapture(camera_path)
try:
first_loop = True
while (capture.isOpened()):
ret, img = capture.read()
if img is None:
continue
# Initialize video writing
if (first_loop):
first_loop = False
fourcc = cv2.VideoWriter_fourcc(*'H264')
h, w = img.shape[:2]
writer = cv2.VideoWriter('test.mkv', fourcc, 10, (w, h), True)
t1 = time.time()
boxes, boxes_c = mtcnn_detector.detect_pnet(img)
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
print('shape: ', img.shape, '--', 'time: ', time.time() - t1)
draw = img.copy()
if boxes_c is not None:
font = cv2.FONT_HERSHEY_SIMPLEX
for b in boxes_c:
cv2.rectangle(draw, (int(b[0]), int(b[1])),
(int(b[2]), int(b[3])), (0, 255, 255), 1)
cv2.putText(draw, '%.3f' % b[4], (int(b[0]), int(b[1])),
font, 0.4, (255, 255, 255), 1)
writer.write(draw)
except KeyboardInterrupt as e:
print("KeyboardInterrupt")
writer.release()
def test_net(prefix,
epoch,
batch_size,
ctx,
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window: # 使用滑动窗口(MTCNN的P_NET不使用了滑动窗口,而是全卷积网络)
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[1], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
# img = cv2.imread('test01.jpg') # 读取图片
# img = cv2.imread('000007.jpg') # 读取图片
# img = cv2.imread('crow.jpg') # 读取图片
img = cv2.imread('physics.jpg') # 读取图片
t1 = time.time() # 开始计时
boxes, boxes_c = mtcnn_detector.detect_pnet(img)
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
print(boxes_c)
print('time: ', time.time() - t1)
if boxes_c is not None:
draw = img.copy()
font = cv2.FONT_HERSHEY_SIMPLEX # Python 一种字体
for b in boxes_c: # nms和bbr之后的结果
# for b in boxes: # nms和bbr之前的结果
# 在draw上绘制矩形框(左上角坐标+右下角坐标)
cv2.rectangle(draw, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])),
(0, 255, 255), 1)
# 在draw上添加文字
cv2.putText(draw, '%.3f' % b[4], (int(b[0]), int(b[1])), font, 0.4,
(255, 255, 255), 1)
cv2.imshow("detection result", draw)
cv2.waitKey(0)
def test_net(dataset_path,
prefix,
faceRecognize_model,
epoch,
batch_size,
test_mode="onet",
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
margin=44,
stride=2,
slide_window=False):
detectors = [None, None, None]
model_path = ['%s-%s' % (x, y) for x, y in zip(prefix, epoch)]
if slide_window:
PNet = Detector(P_Net, 12, batch_size[0], model_path[0])
else:
PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet
#load rnet model
if test_mode in ["rnet", "onet"]:
detectors[1] = Detector(R_Net, 24, batch_size[1], model_path[1])
# load onet model
if test_mode == "onet":
detectors[2] = Detector(O_Net, 48, batch_size[2], model_path[2])
# load faceRecognize model
face_rec = FcnRecognize(faceRecognize_model, data_size=67, batch_size=16)
mtcnn_detector = MtcnnDetector(detectors=detectors,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
######load data
input_dir0 = '/mllib/ALG/facenet-tensorflow/jz_80val'
classes1 = os.listdir(input_dir0)
message_path = []
for cls1 in classes1:
classes2_path = os.path.join(input_dir0, cls1)
try:
classes2 = os.listdir(classes2_path)
except Exception as e:
print e
continue
key = cv2.waitKey(1)
if key == 'q': ### when keycode is q
print('====------======\n')
img_list_tmp = []
num_id = 0
image_message = {}
message = []
meg_name = [
'exit_id', 'img_read', 'img_detect', 'face_num', 'face_roi',
'exit_person', 'score', 'person_name'
]
##########0 exit_id :Is there an id ; 1 ,0
##########1 img_read :Whether to read successfully image ; 1 ,0
##########2 img_detect :Whether to detect successfully image ; 1 ,0
##########3 face_num :The face amount ; 0, 1, 2, 3, ...
##########4 face_roi :The face of a coordinate ; (x1,y1,w1,h1),(x2,y2,w2,h2) ...
##########5 exit_person :The person amount ; 0, 1, 2, 3, ...
##########6 score :The person score ; 0, 1, 2, 3, ...
##########7 score :The person name ; name0, name1, name2, name3, ...
for cls2 in classes2:
classes3_path = os.path.join(classes2_path, cls2)
print(classes3_path)
if 1: #classes2_path == "emb_check" :
image_message[meg_name[0]] = 1
else:
num_id = 0
image_message[meg_name[0]] = 0
continue
try:
img = cv2.imread(classes3_path)
except Exception as e:
print e
image_message[meg_name[1]] = 0
continue
#img = cv2.imread(image_name)#('test01.jpg')
#img = cv2.imread(input_test+'.jpg')
#cv2.imshow("img", img)
#cv2.waitKey(0)
image_message[meg_name[0]] = 1
t1 = time.time()
try:
boxes, boxes_c = mtcnn_detector.detect_pnet(img)
except Exception as e:
image_message[meg_name[2]] = 0
print e
print(classes3_path)
continue
image_message[meg_name[2]] = 1
#message.append("img_available")
if boxes_c is None:
image_message[meg_name[3]] = 0
continue
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
if boxes_c is None:
image_message[meg_name[3]] = 0
continue
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
if boxes_c is None:
image_message[meg_name[3]] = 0
continue
image_message[meg_name[3]] = len(boxes_c)
print('box_num:%d', len(boxes_c))
print 'time: ', time.time() - t1
message.append("have_face")
num_box = []
if boxes_c is not None:
img0 = misc.imread(classes3_path)
img_size = np.asarray(img0.shape)[0:2]
nrof_samples = len(boxes_c)
for det in boxes_c:
bb = np.zeros(4, dtype=np.int32)
margin_tmp = ((det[3] - det[1]) - (det[2] - det[0])) / 2
if margin_tmp > 0:
size_tmp = (det[3] - det[1]) * 0
bb[0] = np.maximum(det[0] - margin_tmp - size_tmp, 0)
bb[1] = np.maximum(det[1] - size_tmp, 0)
bb[2] = np.minimum(det[2] + margin_tmp + size_tmp,
img_size[1])
bb[3] = np.minimum(det[3] + size_tmp, img_size[0])
else:
size_tmp = (det[2] - det[0]) * 0
bb[0] = np.maximum(det[0] - size_tmp, 0)
bb[1] = np.maximum(det[1] + margin_tmp - size_tmp, 0)
bb[2] = np.minimum(det[2] + size_tmp, img_size[1])
bb[3] = np.minimum(det[3] - margin_tmp + size_tmp,
img_size[0])
cropped = img0[int(bb[1]):int(bb[3]),
int(bb[0]):int(bb[2]), :]
num_box.append(
"%d,%d,%d,%d" %
(int(bb[1]), int(bb[3]), int(bb[0]), int(bb[2])))
#misc.imsave('hebei/%d.png'%i, cropped)
#cropped=misc.imread('/home/ssd/fb_data/casic_cluster_china67/Jin Jong-oh/Jin Jong-oh_8.png')
#cv2.imshow("cropped", cropped)
#cv2.waitKey(500) ###
aligned = misc.imresize(cropped, (67, 67),
interp='bilinear')
#cv2.imshow("aligned", aligned)
#cv2.waitKey(500) ###
#misc.imsave('hebei/%d.png'%i, aligned)
prewhitened = faceRecognize.prewhiten(aligned)
img_list_tmp.append(prewhitened)
key = cv2.waitKey(1)
image_message[meg_name[4]] = num_box
print(image_message)
if len(img_list_tmp) < 2: #1 :
continue
img_list = [None] * len(img_list_tmp)
for i in range(len(img_list_tmp)):
img_list[i] = img_list_tmp[i]
images = np.stack(img_list)
t_emb = time.time()
emb = face_rec.recognize(images)
np.save("./emb_check.txt", emb)
print 'emb_1_time: ', time.time() - t_emb
for i in range(len(img_list_tmp)):
print('%1d ' % i)
score_tmp = 65
score2person = -1
for j in range(len(img_list_tmp)):
dist = np.sqrt(
np.sum(np.square(np.subtract(emb[i, :], emb[j, :]))))
print(' %1.4f ' % dict2score(dist))
print('')
if dict2score(dist) > score_tmp:
score_tmp = dict2score(dist)
score2person = j
if score_tmp > 65:
image_message[meg_name[6]] = score_tmp
image_message[meg_name[7]] = 'num_0'
else:
image_message[meg_name[6]] = 0
image_message[meg_name[7]] = 'num_0'
print('-----====end compare !!===-----')
print(image_message)
# Gathers the input arguments
args = get_arguments()
# Gathering variables from arguments
detector = args.detector
classifier = args.classifier
device = args.device
height = args.height
width = args.width
threshold = args.threshold
# Starts a thread from the `Stream` class
v = Stream(height, width, device).start_thread()
# Loading the detection model
det = Detector.load(f'models/{detector}')
# Loading the classification model
clf = Classifier.load(f'models/{classifier}')
# While the loop is True
while True:
# Reads a new frame
valid, frame = v.read()
# Checks if the frame is valid
if valid:
# Performing the detection over the frame
det_preds = det(frame)
# Detects bounding boxes over the objects
def test_net(prefix,
epoch,
batch_size,
ctx,
thresh=[0.6, 0.6, 0.7],
min_face_size=24,
stride=2,
slide_window=False,
camera_path='0'):
detectors = [None, None, None]
# load pnet model
args, auxs = load_param(prefix[0], epoch[0], convert=True, ctx=ctx)
if slide_window:
PNet = Detector(P_Net("test"), 12, batch_size[0], ctx, args, auxs)
else:
PNet = FcnDetector(P_Net("test"), ctx, args, auxs)
detectors[0] = PNet
# load rnet model
args, auxs = load_param(prefix[1], epoch[0], convert=True, ctx=ctx)
RNet = Detector(R_Net("test"), 24, batch_size[1], ctx, args, auxs)
detectors[1] = RNet
# load onet model
args, auxs = load_param(prefix[2], epoch[2], convert=True, ctx=ctx)
ONet = Detector(O_Net("test"), 48, batch_size[2], ctx, args, auxs)
detectors[2] = ONet
mtcnn_detector = MtcnnDetector(detectors=detectors,
ctx=ctx,
min_face_size=min_face_size,
stride=stride,
threshold=thresh,
slide_window=slide_window)
try:
capture = cv2.VideoCapture(int(camera_path))
except ValueError as e:
capture = cv2.VideoCapture(camera_path)
while (capture.isOpened()):
ret, img = capture.read()
if img is None:
continue
# img = cv2.imread('test01.jpg')
t1 = time.time()
boxes, boxes_c = mtcnn_detector.detect_pnet(img)
boxes, boxes_c = mtcnn_detector.detect_rnet(img, boxes_c)
boxes, boxes_c = mtcnn_detector.detect_onet(img, boxes_c)
print('shape: ', img.shape, '--', 'time: ', time.time() - t1)
if boxes_c is not None:
draw = img.copy()
font = cv2.FONT_HERSHEY_SIMPLEX
for b in boxes_c:
cv2.rectangle(draw, (int(b[0]), int(b[1])),
(int(b[2]), int(b[3])), (0, 255, 255), 1)
cv2.putText(draw, '%.3f' % b[4], (int(b[0]), int(b[1])), font,
0.4, (255, 255, 255), 1)
cv2.imshow("detection result", draw)
else:
cv2.imshow("detection result", img)
k = cv2.waitKey(1)
if k == 27 or k == 113: # Esc or q key to stop
break
class Trainer:
def __init__(self, classes_info, model_info, size):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.det = Detector(classes_info, model_info, self.device)
# self.det.load_model("checkpoints/model_0.6161.pth")
self.has_train_config = False
self.size = size
def make_train_config(self, image_path, train_path, val_path, lr,
batch_size, num_workers):
self._make_optimizer(lr)
self._make_dataset(image_path, train_path, batch_size, val_path,
num_workers)
self._make_criterion()
self.min_loss = float("inf")
self.has_train_config = True
def _make_dataset(self,
image_path,
train_path,
batch_size,
val_path=None,
num_workers=6):
train_trans = [
CropResize(self.size),
Normalization([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
]
train_trans = Compose(train_trans)
train_data = MaskData(image_path, train_path, (self.size, self.size),
train_trans)
self.train_loader = DataLoader(train_data,
batch_size,
True,
num_workers=num_workers,
pin_memory=True)
if val_path:
val_trans = [
CropResize(self.size),
Normalization([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
]
val_trans = Compose(val_trans)
val_data = MaskData(image_path, val_path, (self.size, self.size),
val_trans)
self.val_loader = DataLoader(val_data,
batch_size,
num_workers=num_workers,
pin_memory=True)
else:
self.val_loader = None
def _make_criterion(self):
self.criterion = FocalLoss()
def _make_optimizer(self, lr):
self.op = Adam(self.det.get_param(), lr=lr, weight_decay=5e-4)
def fit_one_epoch(self):
assert self.has_train_config, "还没配置好训练参数"
avg_loss = AVGMetrics("train_loss")
self.det.set_status("train")
with tqdm(self.train_loader, desc="train") as pbar:
for data in pbar:
self.op.zero_grad()
data = {k: v.to(self.device) for k, v in data.items()}
images = data["image"]
pred = self.det.inference(images)
loss, scalar = self.criterion(pred, data)
avg_loss.update(loss.item(), len(images))
loss.backward()
self.op.step()
pbar.set_postfix(**scalar)
return str(avg_loss)
def eval_one_epoch(self, save_dir, prefix=""):
assert self.val_loader, "验证集没有构建"
avg_loss = AVGMetrics("val_loss")
self.det.set_status("eval")
with torch.no_grad():
with tqdm(self.train_loader, desc="eval") as pbar:
for data in pbar:
data = {k: v.to(self.device) for k, v in data.items()}
images = data["image"]
pred = self.det.inference(images)
loss, scalar = self.criterion(pred, data)
avg_loss.update(loss.item(), len(images))
pbar.set_postfix(**scalar)
if self.min_loss > avg_loss():
self.min_loss = avg_loss()
self.det.save_model(f"{save_dir}/{prefix}_{self.min_loss}.pth")
print(f"min_loss update to {self.min_loss}")
return str(avg_loss)
def __init__(self, classes_info, model_info, size):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.det = Detector(classes_info, model_info, self.device)
# self.det.load_model("checkpoints/model_0.6161.pth")
self.has_train_config = False
self.size = size
