前言
我想使用一些已有框架、模型，对我们拍好的视频进行推理/预测，得到coco格式的标记框。不久前我使用mmdetection框架进行推理，发现在此框架下可以使用fasterrcnn进行推理，若使用maskrcnn，则会报这样一个错误：
test_cfg specified both in outer and xxx(忘了是啥了) field
摸索了好久也没找到原因。问了gpt说还有一个框架叫detectron2，于是我直接搬起我的小板凳就来到了这里。
我的一些配置（仅供参考）
ubuntu x86_64
cuda11.1
python3.8
torch1.10.0+cu111
torchvision0.11.0+cu111
detectron20.6+cu111
推理
先放完整代码，有大量注释，后续也有部分对代码的解释和模型权重的下载地址。
import numpy as np
import os
import json
import cv2
import time
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from pycocotools import mask as maskUtils
workdir = "/home/xxxxxx/detectron2-main/configs/COCO-InstanceSegmentation/"
model_cfg = "mask_rcnn_R_50_FPN_3x.yaml"
weights_cfg = "model_final_fpn50.pkl"
image_dir = "/home/xxxxxx/inference/image_video/"
image = "oip.jpg"
video = "6.mp4"
coco_before_partion_dir = "/home/xxxxxx/inference/coco_before_partion/"
coco_json = "video.json"
frame_width = 0
frame_height = 0
def main():
# image_path = os.path.join(image_dir, image) # 图片地址/路径
video_path = os.path.join(image_dir, video) # 视频地址=视频目录+视频文件名
model_cfg_path = os.path.join(workdir, model_cfg) # 模型配置参数地址/路径
weights_cfg_path = os.path.join(workdir, weights_cfg) # 权重配置地址/路径
threshold = 0.8 # 置信度阈值
# 生成的coco数据集的地址（目录），没有划分训练集和测试集
# image_save_path = os.path.join(coco_before_partion_dir, coco_json)
video_save_path = os.path.join(coco_before_partion_dir, coco_json)
# 1.
predictor, metadata = load_model(
model_cfg_path = model_cfg_path,
weights_cfg_path = weights_cfg_path,
threshold = threshold
)
# 推理视频，将每一帧的推理结果一起放在一个json文件里
results = inference_video(video_path, predictor)
# 制作coco数据集
make_coco_video(results, video_save_path, metadata)
def load_model(model_cfg_path, weights_cfg_path, threshold):
print("-------------------------------------")
print("开始加载模型...")
# 记录开始时间
t = time.time()
# 初始化配置变量
cfg = get_cfg()
# 加载模型配置文件
cfg.merge_from_file(model_cfg_path)
# 加载预训练权重文件
cfg.MODEL.WEIGHTS = weights_cfg_path
# 设置置信度阈值，置信度小于阈值的框会被丢弃，只有大于等于它才能被保留
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
# 生成一个预测器，一会对于视频的每一帧都用这个预测器进行预测/推理
predictor = DefaultPredictor(cfg)
# 包含模型信息的元数据（metadata）对象。该元数据提供了关于已加载模型的各种信息
metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
# 计算耗时
print("加载模型完成!\n耗时{} s".format(time.time() - t))
print("-------------------------------------")
return predictor, metadata
def inference_video(video_path, predictor):
# 读取要推理的视频
cap = cv2.VideoCapture(video_path)
# 得到视频的宽、高以及总帧数
global frame_width, frame_height
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print("-------------------------------------")
print("total frames:", total_frames)
print("frame width:", frame_width)
print("frame height:", frame_height)
print("-------------------------------------")
# 开始推理/预测
print("开始推理...")
# 定义一个空的列表用于存储结果
results = []
# 设置初始图片image、标注annotations的id
image_id = 1
annotations_id = 1
# 创建进度条对象
with tqdm(total = total_frames, desc="预测/推理视频帧", unit=" frames") as pbar:
while cap.isOpened():
# 读取视频帧
ret, frame = cap.read()
if not ret:
break
# 转换帧格式，BGR to RGB
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# 一帧图片的预测结果
outputs = predictor(frame)
# 获取预测结果.
"""
用于将Instances对象从GPU转移到CPU上。这样做的目的是为了能够使用常规的CPU操作和库(例如numpy和OpenCV)处理和可视化结果。
"""
instances = outputs["instances"].to("cpu")
# 遍历每一个实例
# # 先将预测结果的矩形框转成一个列表，一会直接根据索引使用其中的元素
bbox_list = instances.pred_boxes.tensor.tolist()
for i in range(len(instances)):
# 获取实例的类别
class_id = instances.pred_classes[i]
# 获取实例的置信度
score = instances.scores[i]
# 获取实例的掩码
mask = instances.pred_masks[i]
# 获取实例的边界框
# bbox = instances.pred_boxes[i]
# 使用 RLE 编码格式将掩码图像转换为分割区域表示
segmentation = maskUtils.encode(np.asfortranarray(mask,dtype=np.uint8))
segmentation["counts"] = segmentation["counts"].decode("utf-8")
# 将结果添加到列表中
result = {
"id": annotations_id,
"image_id": image_id,
"category_id": class_id.item(),
"score": score.item(),
"segmentation": segmentation,
"bbox": bbox_list[i], # 一个实例一个框
"area": int(maskUtils.area(segmentation))
}
results.append(result)
annotations_id += 1
image_id += 1 # 增加图像ID计数器
# 进度条
pbar.update(1)
# 释放视频资源
cap.release()
cv2.destroyAllWindows()
print("推理完成!")
print("视频包含实例个数:{}".format(len(results)))
print("-------------------------------------")
return results
def make_coco_video(results, save_path, metadata):
print("-------------------------------------")
print("正在生成coco格式数据集")
print("-------------------------------------")
# 定义一个字典用于存储最终结果
coco_output = {
"images": [],
"annotations": [],
"categories": []
}
# 用于记录图像ID的字典。
'''
因为我们上一步预测的结果（一个列表），这个列表的每一个元素并不是一帧图片的所有实例，
而是每个实例，意思就是每个实例是列表的一个元素。一帧图片可能有多个实例，也可能只有一个实例。
所有实例是平等的，都在列表中占一个元素。
后续我们需要将所有的图片id存起来，不能有冗余，但是我们提取的时候是遍历每个实例，
从其中的'image_id'字段提取id，直接提取会有冗余，可以设置一个字典，每添加一个image_id就往
里添加一个元素，再添加下一个实例的image_id时，检查这个id是否已经存在，不存在则添加，反之则不添加
'''
# 设置image_id字典
image_id_dict = {}
# 创建一个进度条对象
with tqdm(total=len(results), desc="转成coco格式", unit=" instances") as pbar:
# 遍历每个实例，将结果按图像进行分组
for i, result in enumerate(results):
image_id = result["image_id"]
# 检查图像ID是否已经在字典中存在
if image_id not in image_id_dict:
# 添加图像信息
image_info = {
"id": image_id,
"file_name": "frame{}.jpg".format(image_id),
"width": frame_width,
"height": frame_height
}
# 1.添加image信息
coco_output["images"].append(image_info)
# 添加图像ID到字典中
image_id_dict[image_id] = len(coco_output["images"]) - 1
# 2.添加注释信息
annotation_info = {
"id": result["id"],
"image_id": result["image_id"],
"category_id": result["category_id"],
"segmentation": result["segmentation"],
"area": result["area"],
"bbox": result["bbox"],
"iscrowd": 0,
"score": result["score"]
}
coco_output["annotations"].append(annotation_info)
pbar.update(1)
# 3.更新"categories"字段
for class_id, class_name in enumerate(metadata.get("thing_classes", [])):
category_info = {
"id": class_id,
"name": class_name,
"supercategory": ""
}
coco_output["categories"].append(category_info)
print("-------------------------------------")
print("生成coco数据集成功!")
print("-------------------------------------")
# 将最终结果保存为JSON文件
with open(save_path, "w") as f:
json.dump(coco_output, f)
if __name__ == '__main__':
main()
0.先设置一些全局变量：
workdir = "/home/yhy/jzh/detectron2-main/configs/COCO-InstanceSegmentation/"
model_cfg = "mask_rcnn_R_50_FPN_3x.yaml"
weights_cfg = "model_final_fpn50.pkl"
image_dir = "/home/yhy/jzh/inference/image_video/"
image = "oip.jpg"
video = "6.mp4"
coco_before_partion_dir = "/home/yhy/jzh/inference/coco_before_partion/"
coco_json = "video.json"
- 工作空间目录
workdir - 模型参数
model_cfg模型参数一般在本地都有，在detectron2的configs目录里，里面有很多yaml文件，对应不同模型的参数 - 权重文件
weights_cfg这个需要去官网下载：https://github.com/facebookresearch/detectron2/blob/main/MODEL_ZOO.md
右侧的download下载地址
- image_dir是存放图片/视频的目录，image、video分别是图片文件，视频文件，如果你只推理视频就不需要设置image
- frame_width、frame_height是图片的宽和高，因为一会要在不同的函数中使用，所以设置全局变量（方法感觉不太好，暂时也没想到更好的方法。
main()函数
def main():
# 0.设置一些
# image_path = os.path.join(image_dir, image) # 图片地址/路径
video_path = os.path.join(image_dir, video) # 视频地址=视频目录+视频文件名
model_cfg_path = os.path.join(workdir, model_cfg) # 模型配置参数地址/路径
weights_cfg_path = os.path.join(workdir, weights_cfg) # 权重配置地址/路径
threshold = 0.8 # 置信度阈值
# 生成的coco数据集的地址（目录），没有划分训练集和测试集
# image_save_path = os.path.join(coco_before_partion_dir, coco_json)
video_save_path = os.path.join(coco_before_partion_dir, coco_json)
# 1.加载模型，进行预测/推理
predictor, metadata = load_model(
model_cfg_path = model_cfg_path,
weights_cfg_path = weights_cfg_path,
threshold = threshold
)
# 推理视频，将每一帧的推理结果一起放在一个json文件里
results = inference_video(video_path, predictor)
# 制作coco数据集
make_coco_video(results, video_save_path, metadata)
- 对于load_model函数：
在Detectron2框架中，load_model函数用于加载训练好的模型，并返回一个包含模型信息的元数据（metadata）对象。该元数据提供了关于已加载模型的各种信息，以帮助用户理解和操作模型。
返回值metadata是一个MetadataCatalog对象，包含以下常见属性：
“model_type”：模型类型的字符串表示。例如，对于Mask R-CNN模型，值为"mask_rcnn"。
“thing_classes”：包含检测/分割任务中的类别名称的列表。每个类别都由一个字符串表示。
“stuff_classes”：对于分割任务中的背景类别，包含类别名称的列表。通常在实例分割中使用，与"thing_classes"不同。
“keypoint_names”：如果模型支持关键点检测，则包含关键点名称的列表。
“keypoint_flip_map”：关键点镜像映射表，用于在镜像图像上翻转关键点。
“thing_dataset_id_to_contiguous_id”：将数据集中类别ID映射到模型输出中连续的类别ID的字典。
“stuff_dataset_id_to_contiguous_id”：将数据集中背景类别ID映射到模型输出中连续的类别ID的字典。
“thing_colors”：每个实例类别的颜色编码的列表。这些颜色用于可视化实例分割结果。
“stuff_colors”：每个背景类别的颜色编码的列表。这些颜色用于可视化分割结果中的背景。
我们在后续代码中主要使用其thing_classes属性，生成coco数据集的category_id属性
coco数据格式：
coco_data{['image'], ['annotations'], ['category_id']}
- 推理/预测
results = inference_video(video_path, predictor)
对video_path地址的视频，使用已经加载好的模型预测器进行推理/预测。
- 把结果制作成coco格式数据
make_coco_video(results, video_save_path, metadata)
results是上一步预测的结果，video_save_path是最后要保存的json文件的地址，metadata是我们要填充coco的category_id字段时要用的变量。
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
