安装环境
!pip install openmim
! mim install mmdet
下载模型
!mim download mmdet --config mask_rcnn_x101_32x4d_fpn_2x_coco
修改 配置文件,使用继承的方式
mim train mmdet balloon.py
balloon.py
_base_ = ['mask_rcnn_x101_32x4d_fpn_2x_coco.py']
data = dict(
train=dict(
ann_file='/openbayes/home/train_data.json',
img_prefix='/openbayes/input/input0/balloon/train',
classes = ("balloon",)
),
val=dict(
ann_file='/openbayes/home/val_data.json',
img_prefix='/openbayes/input/input0/balloon/val',
classes = ("balloon",)
),
test=dict(
ann_file='/openbayes/home/val_data.json',
img_prefix='/openbayes/input/input0/balloon/val',
classes = ("balloon",)
)
)
model = dict(
roi_head=dict(
bbox_head=dict(
num_classes=1
),
mask_head=dict(
num_classes=1
)
)
)
python data2cocos ann_file, out_file, image_prefix
ann_file 标注文件
Outfile 保存文件
image_prefix image路径
data2cocos.py
import os.path as osp
import sys, mmcv
def convert_balloon_to_coco(ann_file, out_file, image_prefix):
data_infos = mmcv.load(ann_file)
annotations = []
images = []
obj_count = 0
for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())):
filename = v['filename']
img_path = osp.join(image_prefix, filename)
height, width = mmcv.imread(img_path).shape[:2]
images.append(dict(
id=idx,
file_name=filename,
height=height,
width=width))
bboxes = []
labels = []
masks = []
for _, obj in v['regions'].items():
assert not obj['region_attributes']
obj = obj['shape_attributes']
px = obj['all_points_x']
py = obj['all_points_y']
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
poly = [p for x in poly for p in x]
x_min, y_min, x_max, y_max = (
min(px), min(py), max(px), max(py))
data_anno = dict(
image_id=idx,
id=obj_count,
category_id=0,
bbox=[x_min, y_min, x_max - x_min, y_max - y_min],
area=(x_max - x_min) * (y_max - y_min),
segmentation=[poly],
iscrowd=0)
annotations.append(data_anno)
obj_count += 1
coco_format_json = dict(
images=images,
annotations=annotations,
categories=[{'id':0, 'name': 'balloon'}])
mmcv.dump(coco_format_json, out_file)
if __name__ == '__main__':
ann_file = sys.argv[1]
out_file = sys.argv[2]
image_prefix = sys.argv[3]
convert_balloon_to_coco(ann_file, out_file, image_prefix)
python mask_video.py test_video.mp4 /openbayes/home/video.mp4 /openbayes/home/work_dirs/balloon/epoch_20.pth balloon.py 5
mask_video.py
import os, sys, math, cv2, mmcv, numpy
from mmdet.apis import inference_detector, init_detector
def mask_draw(src_file, dst_file, mod_file, conf_file, threshold, device):
model = init_detector(conf_file, mod_file, device=device)
score_thr = threshold
frames = mmcv.VideoReader(src_file)
fourcc = "XVID"
fps = math.ceil(frames.fps)
resolution = (frames.width, frames.height)
vwriter = cv2.VideoWriter(dst_file, cv2.VideoWriter_fourcc(*fourcc), fps, resolution)
for fid in mmcv.track_iter_progress(range(len(frames))):
if frames[fid] is None:
continue
result = inference_detector(model, frames[fid])
if len(result[0][0]) == 0:
continue
mask = numpy.zeros(result[1][0][0].shape, dtype=bool)
for bid in range(len(result[0][0])):
if result[0][0][bid][4] < score_thr:
continue
mask = numpy.add(mask, result[1][0][bid])
img = frames[fid].copy()
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
for i in range(resolution[1]):
for j in range(resolution[0]):
if not mask[i][j]:
img[i][j] = [gray[i][j]]*3
vwriter.write(img)
vwriter.release()
if __name__ == '__main__':
src_file = sys.argv[1]
dst_file = sys.argv[2]
mod_file = sys.argv[3]
conf_file = sys.argv[4]
threshold = float(sys.argv[5])
device="cuda:0"
if len(sys.argv) > 6:
device = sys.argv[6]
mask_draw(src_file, dst_file, mod_file, conf_file, threshold, device)