tf-centernet

使用Tensorflow 2 实现 CenterNet。

安装

pip instal tf-centernet

例子

物体检测

import numpy as np
import PIL.Image
import centernet

# Default: num_classes=80
obj = centernet.ObjectDetection(num_classes=80)

# Default: weights_path=None
# num_classes=80 and weights_path=None: Pre-trained COCO model will be loaded.
# Otherwise: User-defined weight file will be loaded.
obj.load_weights(weights_path=None)

img = np.array(PIL.Image.open('./data/sf.jpg'))[..., ::-1]

# The image with predicted bounding-boxes is created if `debug=True`
boxes, classes, scores = obj.predict(img, debug=True)

姿势估计

import numpy as np
import PIL.Image
import centernet

# Default: num_joints=17
pe = centernet.PoseEstimation(num_joints=17)

# Default: weights_path=None
# num_joints=17 and weights_path=None: Pre-trained COCO model will be loaded.
# Otherwise: User-defined weight file will be loaded.
pe.load_weights(weights_path=None)

# Adjust this for the better prediction
pe.score_threshold = 0.1

img = np.array(PIL.Image.open('./data/chi.jpg'))[..., ::-1]

# The image with predicted keypoints is created if `debug=True`
boxes, keypoints, scores = pe.predict(img, debug=True)

去做

• 物体检测
• 使用 Hourglass-104 进行目标检测的预训练模型
• 姿势估计
• 使用 Hourglass-104 进行姿势估计的预训练模型
• DLA-34 主干和预训练模型
• 训练函数和损失定义
• 训练数据增强