def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
#----------------------------------------#
# 计算种类数量
#----------------------------------------#
self.num_classes = len(self.class_names)
#----------------------------------------#
# 创建Efficientdet模型
#----------------------------------------#
self.Efficientdet = Efficientdet(self.phi, self.num_classes)
self.Efficientdet.load_weights(self.model_path,
by_name=True,
skip_mismatch=True)
print('{} model, anchors, and classes loaded.'.format(model_path))
# 画框设置不同的颜色
hsv_tuples = [(x / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
self.colors))
class EfficientDet(object):
_defaults = {
"model_path": 'model_data/efficientdet-d0-voc.h5',
"classes_path": 'model_data/voc_classes.txt',
"phi": 0,
"confidence": 0.4,
"iou": 0.3,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化efficientdet
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
self.class_names = self._get_class()
self.model_image_size = [
image_sizes[self.phi], image_sizes[self.phi], 3
]
self.sess = K.get_session()
self.generate()
self.bbox_util = BBoxUtility(self.num_classes, nms_thresh=self.iou)
self.prior = self._get_prior()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def _get_class(self):
classes_path = os.path.expanduser(self.classes_path)
with open(classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
#---------------------------------------------------#
# 获得先验框
#---------------------------------------------------#
def _get_prior(self):
data = get_anchors(image_sizes[self.phi])
return data
#---------------------------------------------------#
# 载入模型
#---------------------------------------------------#
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
#----------------------------------------#
# 计算种类数量
#----------------------------------------#
self.num_classes = len(self.class_names)
#----------------------------------------#
# 创建Efficientdet模型
#----------------------------------------#
self.Efficientdet = Efficientdet(self.phi, self.num_classes)
self.Efficientdet.load_weights(self.model_path)
print('{} model, anchors, and classes loaded.'.format(model_path))
# 画框设置不同的颜色
hsv_tuples = [(x / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
self.colors))
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
#---------------------------------------------------------#
image = image.convert('RGB')
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
#---------------------------------------------------------#
crop_img = letterbox_image(
image, [self.model_image_size[1], self.model_image_size[0]])
#-----------------------------------------------------------#
# 图片预处理，归一化。获得的photo的shape为[1, 512, 512, 3]
#-----------------------------------------------------------#
photo = np.array(crop_img, dtype=np.float32)
photo = np.reshape(preprocess_input(photo), [
1, self.model_image_size[0], self.model_image_size[1],
self.model_image_size[2]
])
preds = self.Efficientdet.predict(photo)
#-----------------------------------------------------------#
# 将预测结果进行解码
#-----------------------------------------------------------#
results = self.bbox_util.detection_out(
preds, self.prior, confidence_threshold=self.confidence)
#--------------------------------------#
# 如果没有检测到物体，则返回原图
#--------------------------------------#
if len(results[0]) <= 0:
return image
results = np.array(results)
det_label = results[0][:, 5]
det_conf = results[0][:, 4]
det_xmin, det_ymin, det_xmax, det_ymax = results[0][:, 0], results[
0][:, 1], results[0][:, 2], results[0][:, 3]
#-----------------------------------------------------------#
# 筛选出其中得分高于confidence的框
#-----------------------------------------------------------#
top_indices = [
i for i, conf in enumerate(det_conf) if conf >= self.confidence
]
top_conf = det_conf[top_indices]
top_label_indices = det_label[top_indices].tolist()
top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(
det_xmin[top_indices],
-1), np.expand_dims(det_ymin[top_indices], -1), np.expand_dims(
det_xmax[top_indices],
-1), np.expand_dims(det_ymax[top_indices], -1)
#-----------------------------------------------------------#
# 去掉灰条部分
#-----------------------------------------------------------#
boxes = efficientdet_correct_boxes(
top_ymin, top_xmin, top_ymax, top_xmax,
np.array([self.model_image_size[0], self.model_image_size[1]]),
image_shape)
font = ImageFont.truetype(font='model_data/simhei.ttf',
size=np.floor(3e-2 * np.shape(image)[1] +
0.5).astype('int32'))
thickness = max((np.shape(image)[0] + np.shape(image)[1]) //
self.model_image_size[0], 1)
for i, c in enumerate(top_label_indices):
predicted_class = self.class_names[int(c)]
score = top_conf[i]
top, left, bottom, right = boxes[i]
top = top - 5
left = left - 5
bottom = bottom + 5
right = right + 5
top = max(0, np.floor(top + 0.5).astype('int32'))
left = max(0, np.floor(left + 0.5).astype('int32'))
bottom = min(
np.shape(image)[0],
np.floor(bottom + 0.5).astype('int32'))
right = min(
np.shape(image)[1],
np.floor(right + 0.5).astype('int32'))
# 画框框
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
label = label.encode('utf-8')
print(label, top, left, bottom, right)
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
for i in range(thickness):
draw.rectangle([left + i, top + i, right - i, bottom - i],
outline=self.colors[int(c)])
draw.rectangle(
[tuple(text_origin),
tuple(text_origin + label_size)],
fill=self.colors[int(c)])
draw.text(text_origin,
str(label, 'UTF-8'),
fill=(0, 0, 0),
font=font)
del draw
return image
def get_FPS(self, image, test_interval):
image_shape = np.array(np.shape(image)[0:2])
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
#---------------------------------------------------------#
crop_img = letterbox_image(
image, [self.model_image_size[1], self.model_image_size[0]])
#-----------------------------------------------------------#
# 图片预处理，归一化。获得的photo的shape为[1, 512, 512, 3]
#-----------------------------------------------------------#
photo = np.array(crop_img, dtype=np.float32)
photo = np.reshape(preprocess_input(photo), [
1, self.model_image_size[0], self.model_image_size[1],
self.model_image_size[2]
])
preds = self.Efficientdet.predict(photo)
#-----------------------------------------------------------#
# 将预测结果进行解码
#-----------------------------------------------------------#
results = self.bbox_util.detection_out(
preds, self.prior, confidence_threshold=self.confidence)
if len(results[0]) > 0:
results = np.array(results)
det_label = results[0][:, 5]
det_conf = results[0][:, 4]
det_xmin, det_ymin, det_xmax, det_ymax = results[0][:, 0], results[
0][:, 1], results[0][:, 2], results[0][:, 3]
#-----------------------------------------------------------#
# 筛选出其中得分高于confidence的框
#-----------------------------------------------------------#
top_indices = [
i for i, conf in enumerate(det_conf) if conf >= self.confidence
]
top_conf = det_conf[top_indices]
top_label_indices = det_label[top_indices].tolist()
top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(
det_xmin[top_indices],
-1), np.expand_dims(det_ymin[top_indices], -1), np.expand_dims(
det_xmax[top_indices],
-1), np.expand_dims(det_ymax[top_indices], -1)
#-----------------------------------------------------------#
# 去掉灰条部分
#-----------------------------------------------------------#
boxes = efficientdet_correct_boxes(
top_ymin, top_xmin, top_ymax, top_xmax,
np.array([self.model_image_size[0], self.model_image_size[1]]),
image_shape)
t1 = time.time()
for _ in range(test_interval):
preds = self.Efficientdet.predict(photo)
#-----------------------------------------------------------#
# 将预测结果进行解码
#-----------------------------------------------------------#
results = self.bbox_util.detection_out(
preds, self.prior, confidence_threshold=self.confidence)
if len(results[0]) > 0:
results = np.array(results)
det_label = results[0][:, 5]
det_conf = results[0][:, 4]
det_xmin, det_ymin, det_xmax, det_ymax = results[
0][:, 0], results[0][:, 1], results[0][:, 2], results[0][:,
3]
#-----------------------------------------------------------#
# 筛选出其中得分高于confidence的框
#-----------------------------------------------------------#
top_indices = [
i for i, conf in enumerate(det_conf)
if conf >= self.confidence
]
top_conf = det_conf[top_indices]
top_label_indices = det_label[top_indices].tolist()
top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(
det_xmin[top_indices], -1), np.expand_dims(
det_ymin[top_indices], -1), np.expand_dims(
det_xmax[top_indices],
-1), np.expand_dims(det_ymax[top_indices], -1)
#-----------------------------------------------------------#
# 去掉灰条部分
#-----------------------------------------------------------#
boxes = efficientdet_correct_boxes(
top_ymin, top_xmin, top_ymax, top_xmax,
np.array(
[self.model_image_size[0], self.model_image_size[1]]),
image_shape)
t2 = time.time()
tact_time = (t2 - t1) / test_interval
return tact_time
def close_session(self):
self.sess.close()
#--------------------------------------------#
# 该部分代码只用于看网络结构，并非测试代码
# map测试请看get_dr_txt.py、get_gt_txt.py
# 和get_map.py
#--------------------------------------------#
from nets.efficientdet import Efficientdet
model = Efficientdet(0)
model.summary()
for i,layer in enumerate(model.layers):
print(i,layer.name)
#-------------------------------------------#
# 训练前，请指定好phi和model_path
# 二者所使用Efficientdet版本要相同
#-------------------------------------------#
phi = 0
annotation_path = '2007_train.txt'
classes_path = 'model_data/new_classes.txt'
class_names = get_classes(classes_path)
NUM_CLASSES = len(class_names)
#-------------------------------------------#
# 权值文件的下载请看README
#-------------------------------------------#
model_path = "model_data/efficientdet-d0-voc.h5"
model = Efficientdet(phi, num_classes=NUM_CLASSES)
priors = get_anchors(image_sizes[phi])
bbox_util = BBoxUtility(NUM_CLASSES, priors)
model.load_weights(model_path, by_name=True, skip_mismatch=True)
# 0.1用于验证，0.9用于训练
val_split = 0.1
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
np.random.seed(None)
num_val = int(len(lines) * val_split)
num_train = len(lines) - num_val
phi = 0
annotation_path = '2007_train.txt'
classes_path = 'model_data/voc_classes.txt'
class_names = get_classes(classes_path)
NUM_CLASSES = len(class_names)
#-------------------------------------------#
# 权值文件的下载请看README
#-------------------------------------------#
model_path = "model_data/efficientdet-d0-voc.h5"
#-------------------------------#
# Dataloder的使用
#-------------------------------#
Use_Data_Loader = True
model = Efficientdet(phi, num_classes=NUM_CLASSES)
priors = get_anchors(image_sizes[phi])
bbox_util = BBoxUtility(NUM_CLASSES, priors)
model.load_weights(model_path, by_name=True, skip_mismatch=True)
# 0.1用于验证，0.9用于训练
val_split = 0.1
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
np.random.seed(None)
num_val = int(len(lines) * val_split)
num_train = len(lines) - num_val
class EfficientDet(object):
_defaults = {
"model_path": 'model_data/efficientdet-d0-voc.h5',
"classes_path": 'model_data/voc_classes.txt',
"phi": 0,
"confidence": 0.4,
"iou": 0.3,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化efficientdet
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
self.class_names = self._get_class()
self.model_image_size = [
image_sizes[self.phi], image_sizes[self.phi], 3
]
self.generate()
self.bbox_util = BBoxUtility(self.num_classes, nms_thresh=self.iou)
self.prior = self._get_prior()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def _get_class(self):
classes_path = os.path.expanduser(self.classes_path)
with open(classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
def _get_prior(self):
data = get_anchors(image_sizes[self.phi])
return data
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
# 计算总的种类
self.num_classes = len(self.class_names)
# 载入模型
self.Efficientdet = Efficientdet(self.phi, self.num_classes)
self.Efficientdet.load_weights(self.model_path,
by_name=True,
skip_mismatch=True)
# self.Efficientdet.summary()
print('{} model, anchors, and classes loaded.'.format(model_path))
# 画框设置不同的颜色
hsv_tuples = [(x / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
self.colors))
@tf.function
def get_pred(self, photo):
preds = self.Efficientdet(photo, training=False)
return preds
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
image_shape = np.array(np.shape(image)[0:2])
crop_img = letterbox_image(
image, [self.model_image_size[0], self.model_image_size[1]])
photo = np.array(crop_img, dtype=np.float32)
# 图片预处理，归一化
photo = np.reshape(preprocess_input(photo), [
1, self.model_image_size[0], self.model_image_size[1],
self.model_image_size[2]
])
preds = self.get_pred(photo)
preds = [pred.numpy() for pred in preds]
# 将预测结果进行解码
results = self.bbox_util.detection_out(
preds, self.prior, confidence_threshold=self.confidence)
if len(results[0]) <= 0:
return image
results = np.array(results)
# 筛选出其中得分高于confidence的框
det_label = results[0][:, 5]
det_conf = results[0][:, 4]
det_xmin, det_ymin, det_xmax, det_ymax = results[0][:, 0], results[
0][:, 1], results[0][:, 2], results[0][:, 3]
top_indices = [
i for i, conf in enumerate(det_conf) if conf >= self.confidence
]
top_conf = det_conf[top_indices]
top_label_indices = det_label[top_indices].tolist()
top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(
det_xmin[top_indices],
-1), np.expand_dims(det_ymin[top_indices], -1), np.expand_dims(
det_xmax[top_indices],
-1), np.expand_dims(det_ymax[top_indices], -1)
# 去掉灰条
boxes = efficientdet_correct_boxes(
top_ymin, top_xmin, top_ymax, top_xmax,
np.array([self.model_image_size[0], self.model_image_size[1]]),
image_shape)
font = ImageFont.truetype(font='model_data/simhei.ttf',
size=np.floor(3e-2 * np.shape(image)[1] +
0.5).astype('int32'))
thickness = (np.shape(image)[0] +
np.shape(image)[1]) // self.model_image_size[0]
for i, c in enumerate(top_label_indices):
predicted_class = self.class_names[int(c)]
score = top_conf[i]
top, left, bottom, right = boxes[i]
top = top - 5
left = left - 5
bottom = bottom + 5
right = right + 5
top = max(0, np.floor(top + 0.5).astype('int32'))
left = max(0, np.floor(left + 0.5).astype('int32'))
bottom = min(
np.shape(image)[0],
np.floor(bottom + 0.5).astype('int32'))
right = min(
np.shape(image)[1],
np.floor(right + 0.5).astype('int32'))
# 画框框
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
label = label.encode('utf-8')
print(label)
if top - label_size[1] >= 0:
text_origin = np.array([left, top - label_size[1]])
else:
text_origin = np.array([left, top + 1])
for i in range(thickness):
draw.rectangle([left + i, top + i, right - i, bottom - i],
outline=self.colors[int(c)])
draw.rectangle(
[tuple(text_origin),
tuple(text_origin + label_size)],
fill=self.colors[int(c)])
draw.text(text_origin,
str(label, 'UTF-8'),
fill=(0, 0, 0),
font=font)
del draw
return image
