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