def main():
args = parse_args()
model = init_detector(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
camera = cv2.VideoCapture(args.camera_id)
print('Press "Esc", "q" or "Q" to exit.')
count = 0
while True:
ret_val, img = camera.read()
prev_time = time.time()
result = inference_detector(model, img)
logger.info('process frame time:' + str(time.time() - prev_time))
demo_path = './results/' + str(count) + '.jpg'
if ret_val:
prev_time = time.time()
show_result_ins(img,
result,
model.CLASSES,
score_thr=0.25,
out_file=demo_path)
logger.info('postprocessing frame time:' +
str(time.time() - prev_time))
else:
break
count += 1
def run_on_onnxruntime():
config_file = '../configs/solov2/solov2_light_448_r34_fpn_8gpu_3x.py'
onnx_file = 'weights/SOLOv2_light_R34.onnx'
input_names = ['input']
# output_names = ['C0', 'C1', 'C2', 'C3', 'C4']
# output_names = ['cate_pred_0', 'cate_pred_1', 'cate_pred_2', 'cate_pred_3', 'cate_pred_4',
# 'kernel_pred_0', 'kernel_pred_1', 'kernel_pred_2', 'kernel_pred_3', 'kernel_pred_4',
# 'seg_pred'] # Origin
output_names = ['cate_pred', 'kernel_pred', 'seg_pred'] # add permute & concate
if isinstance(config_file, str):
cfg = mmcv.Config.fromfile(config_file)
elif not isinstance(config_file, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
'but got {}'.format(type(config_file)))
# 1. Preprocess
# input demo img size 427x640 --> resized 448x671 --> pad 448x672
img = 'images/demo.jpg'
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
# 2. Run inference on onnxruntime
print("Load onnx model from {}.".format(onnx_file))
sess = rt.InferenceSession(onnx_file)
tic = cv2.getTickCount()
onnx_output = sess.run(output_names, {input_names[0]: data['img'][0].unsqueeze(0).cpu().numpy()})
print('-----> onnxruntime inference time: {}ms'.format((cv2.getTickCount() - tic) * 1000/ cv2.getTickFrequency()))
# 3. Get seg
# 调用pytorch定义的获取分割图和matrix nms 以及后处理
from mmdet.models.anchor_heads.solov2_head import SOLOv2Head
solov2_head = SOLOv2Head(num_classes=81,
in_channels=256,
num_grids=[40, 36, 24, 16, 12],
strides=[8, 8, 16, 32, 32],
ins_out_channels = 128,
loss_ins=cfg.model.bbox_head.loss_ins,
loss_cate=cfg.model.bbox_head.loss_cate)
cate_preds = [torch.from_numpy(x) for x in onnx_output[:1]]
kernel_preds = [torch.from_numpy(x) for x in onnx_output[1:2]]
seg_pred = torch.from_numpy(onnx_output[2])
result = solov2_head.get_seg(cate_preds, kernel_preds, seg_pred, [data['img_meta'][0].data], cfg.test_cfg, rescale=True)
show_result_ins(img, result, get_classes('coco'), score_thr=0.25, out_file="images/demo_out_onnxrt_solov2.jpg")
print('Script done!')
def solo_infer(solo_cp_path, src_folder, dst_folder, json_path):
config_file = '../configs/solo/decoupled_solo_r50_fpn_8gpu_3x.py'
# 加载solo模型
model = init_detector(config_file, solo_cp_path, device='cuda:0')
# 创建json文件
fjson = open(json_path, 'w')
# 读取所有输入图像的路径
all_img_paths = glob.glob(os.path.join(src_folder, '*/*.*'))
num_img = len(all_img_paths)
# 开始处理
state = {'num': float(num_img)}
content = [] # 所有图像的信息
for i in tqdm(range(num_img)):
img_path = all_img_paths[i]
img_name = img_path.split('/')[-1]
if cv2.imread(img_path) is None:
print(img_path)
continue
result = inference_detector(model, img_path) # 单张图像通过solo
dst_path = os.path.join(dst_folder, img_name) # 以相同图像名称保存
show_result_ins(img_path,
result,
model.CLASSES,
score_thr=0.25,
out_file=dst_path) # 绘制 mask并保存
# 分析mask信息
img_info = post_treatment(img_path,
result,
model.CLASSES,
score_thr=0.25)
# print(img_info)
content.append(img_info)
# 保存到json文件
state['content'] = content
json.dump(state, fjson, indent=4)
fjson.close()
def main():
args = parse_args()
model = init_detector(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
camera = cv2.VideoCapture(args.camera_id)
print('Press "Esc", "q" or "Q" to exit.')
while True:
ret_val, img = camera.read()
result = inference_detector(model, img)
ch = cv2.waitKey(1)
if ch == 27 or ch == ord('q') or ch == ord('Q'):
break
img_show = show_result_ins(img, result, model.CLASSES, score_thr=0.25)
cv2.imshow('Demo', img_show)
def run_on_tensorrt():
from deploy import common
config_file = '../configs/solov2/solov2_light_448_r34_fpn_8gpu_3x.py'
onnx_file = 'weights/SOLOv2_light_R34.onnx'
input_names = ['input']
# output_names = ['C0', 'C1', 'C2', 'C3', 'C4']
# output_names = ['cate_pred_0', 'cate_pred_1', 'cate_pred_2', 'cate_pred_3', 'cate_pred_4',
# 'kernel_pred_0', 'kernel_pred_1', 'kernel_pred_2', 'kernel_pred_3', 'kernel_pred_4',
# 'seg_pred'] # Origin
output_names = ['cate_pred', 'kernel_pred', 'seg_pred'] # add permute & concate
if isinstance(config_file, str):
cfg = mmcv.Config.fromfile(config_file)
elif not isinstance(config_file, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
'but got {}'.format(type(config_file)))
# 1. Preprocess
# input demo img size 427x640 --> resized 448x671 --> pad 448x672
img = 'images/demo.jpg'
# build the data pipeline
test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = dict(img=img)
data = test_pipeline(data)
# 2. Run inference on trt
print("Load onnx model from {}.".format(onnx_file))
image_shape = data['img_meta'][0].data['pad_shape']
input_shapes = ((1, image_shape[2], image_shape[0], image_shape[1]),) # explict shape
# input_shapes = ((1, 3, 448, 448), (1, 3, 608, 608), (1, 3, 768, 768)) # dynamic shape
# shape_matrix = [
# [1, 40, 40, 80],
# [1, 36, 36, 80],
# [1, 24, 24, 80],
# [1, 16, 16, 80],
# [1, 12, 12, 80],
# [1, 128, 40, 40],
# [1, 128, 36, 36],
# [1, 128, 24, 24],
# [1, 128, 16, 16],
# [1, 128, 12, 12],
# [1, 128, image_shape[0]//4, image_shape[1]//4]
# ]
shape_matrix = [
[3872, 80],
[3872, 128],
[1, 128, image_shape[0] // 4, image_shape[1] // 4]
]
with common.get_engine(onnx_file, onnx_file.replace(".onnx", ".engine"),
input_shapes=input_shapes, force_rebuild=False) \
as engine, engine.create_execution_context() as context:
# Notice: Here we only allocate device memory for speed up
# DYNAMIC shape
# context.active_optimization_profile = 0
# [context.set_binding_shape(x, tuple(y)) for x, y in enumerate(shape_matrix)]
# inputs, outputs, bindings, stream = common.allocate_buffersV2(engine, context)
# EXPLICIT shape
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
# Speed test: cpu(0.976s) vs gpu(0.719s)
# ==> Set host input to the data.
# The common.do_inference function will copy the input to the GPU before executing.
inputs[0].host = data['img'][0].unsqueeze(0).cpu().numpy() # for torch.Tensor
# ==> Or set device input to the data.
# in this mode, common.do_inference function should not copy inputs.host to inputs.device anymore.
# c_type_pointer = ctypes.c_void_p(int(inputs[0].device))
# x.cpu().numpy().copy_to_external(c_type_pointer)
tic = cv2.getTickCount()
trt_outputs = common.do_inferenceV2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream,
batch_size=1, h_=image_shape[0], w_=image_shape[1])
print('-----> tensorRT inference time: {}ms'.format((cv2.getTickCount() - tic) * 1000 / cv2.getTickFrequency()))
# 3. Get seg
# 调用pytorch定义的获取分割图和matrix nms 以及后处理
from mmdet.models.anchor_heads.solov2_head import SOLOv2Head
solov2_head = SOLOv2Head(num_classes=81,
in_channels=256,
num_grids=[40, 36, 24, 16, 12],
strides=[8, 8, 16, 32, 32],
ins_out_channels = 128,
loss_ins=cfg.model.bbox_head.loss_ins,
loss_cate=cfg.model.bbox_head.loss_cate)
# TODO: tensorrt output order is different from pytorch?
# Origin
# ids = [8, 9, 7, 6, 5, 3, 4, 2, 1, 0, 10]
# ids = [9, 8, 7, 5, 6, 4, 3, 2, 0, 1, 10] # TODO: tensorrt output order is different from pytorch?
# Add permute & concate
ids = [1, 0, 2]
cate_preds = [torch.from_numpy(trt_outputs[x]).reshape(y) for x, y in zip(ids[:1], shape_matrix[:1])]
kernel_preds = [torch.from_numpy(trt_outputs[x]).reshape(y) for x, y in zip(ids[1:2], shape_matrix[1:2])]
seg_pred = torch.from_numpy(trt_outputs[2]).reshape(shape_matrix[2])
result = solov2_head.get_seg(cate_preds, kernel_preds, seg_pred, [data['img_meta'][0].data], cfg.test_cfg, rescale=True)
show_result_ins(img, result, get_classes('coco'), score_thr=0.25, out_file="images/demo_out_trt_solov2.jpg")
print('Script done!')
from mmdet.apis import init_detector, inference_detector, show_result_pyplot, show_result_ins
import mmcv
config_file = '../configs/solov2/solov2_x101_dcn_fpn_8gpu_3x.py'
# download the checkpoint from model zoo and put it in `checkpoints/`
checkpoint_file = '../checkpoints/SOLOv2_X101_DCN_3x.pth'
# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# test a single image
img = '/home/baikai/Desktop/AliComp/datasets/PreRoundData/JPEGImages/628058/00001.jpg'
result = inference_detector(model, img)
show_result_ins(img,
result,
model.CLASSES,
score_thr=0.25,
out_file="demo_out.jpg")
config_file = './cfg/solov2_r101.py'
checkpoint_file = '../solo-vos/user_data/model_data/solov2_9cls.pth'
# config_file = './cfg/solov2_light_448_r50.py'
# checkpoint_file = '../solo-vos/user_data/model_data/SOLOv2_LIGHT_448_R50_3x.pth'
model = init_detector(config_file, checkpoint_file, device='cpu')
# # test a single image
# img = '00001.jpg'
# result = inference_detector(model, img)
# show_result_ins(img, result, model.CLASSES, score_thr=0.2,
# out_file='00001.png')
imgs = sorted(glob.glob('/versa/dataset/TIANCHI2021/preTest/*/*.jpg'))[:50]
cost_time = 0
for idx, img_path in enumerate(imgs):
img = cv2.imread(img_path)
name = img_path.split('/')[-1]
video_id = img_path.split('/')[-2]
st = time.time()
result = inference_detector(model, img_path)
cost_time += (time.time() - st)
print(idx, time.time() - st)
save_dir = os.path.join('/versa/dataset/TIANCHI2021/cpu_test', video_id)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
show_result_ins(img, result, model.CLASSES, score_thr=0.2,
out_file=os.path.join(save_dir, name.replace('.jpg', '.png')))
print('Avg cost time per image', cost_time/len(imgs))
def predict(args, model, device):
h, w = 192, 256
k_inv_dot_xy1 = get_coordinate_map(device)
transforms = T.Compose([
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
if os.path.isdir(args.image_path):
images = glob.glob(os.path.join(args.image_path, '*.jpg'))
else:
images = [args.image_path]
with torch.no_grad():
for image_path in images:
image_pil = Image.open(image_path)
raw_w, raw_h = image_pil.size
image = transforms(image_pil)
image = image.to(device).unsqueeze(0)
raw_image = [torch.from_numpy(np.array(image_pil))]
# forward pass
bts_depth, param, mask_lst, instance_map, valid_region, solo_results = \
model(image, raw_image, args.bts_input, args.plane_cls, args.solo_conf)
# # infer per pixel depth using per pixel plane parameter,
# # currently depth_loss need a dummy gt_depth as input
# _, _, per_pixel_depth = depth_loss(param, k_inv_dot_xy1, torch.ones((1, 1, h, w)))
# infer instance depth
# _, plane_depth, _, instance_param = instance_aware_loss(
# mask_lst[0], instance_map[0], param[0], k_inv_dot_xy1,
# torch.ones((1, 1, h, w)), torch.ones((1, 1, h, w)), return_loss=False)
_, _, _, _, plane_depth, instance_param = total_loss2(
mask_lst[0],
param[0],
k_inv_dot_xy1,
valid_region,
torch.ones((1, 1, h, w)),
torch.ones((1, 1, h, w)),
return_loss=False)
plane_depth = plane_depth.cpu().numpy()[0, 0].reshape(h, w)
bts_depth = bts_depth.cpu().numpy()[0, 0] # .reshape(640, 480)
# bts_depth = cv2.resize(bts_depth, (w, h))
valid_region = valid_region[0].cpu().numpy()
# use per pixel depth for non planar region
depth_clean = plane_depth * valid_region
depth = plane_depth * valid_region + bts_depth * (1 - valid_region)
Image.fromarray(depth * 1000).convert('I').resize(
(raw_w, raw_h)).save(image_path.replace('.jpg', '_depth.png'))
Image.fromarray(depth_clean * 1000).convert('I').resize(
(raw_w,
raw_h)).save(image_path.replace('.jpg', '_depth_clean.png'))
# # visualize depth map
# depth = 255 - np.clip(depth / 5 * 255, 0, 255).astype(np.uint8)
# depth = cv2.cvtColor(cv2.resize(depth, (raw_w, raw_h)), cv2.COLOR_GRAY2BGR)
# cv2.imwrite(image_path.replace('.jpg', '.png'), depth)
# visualize solo results
show_result_ins(np.array(image_pil),
solo_results[0],
SOLO_CLASSES,
score_thr=0.2,
out_file=image_path.replace('.jpg', '_blend.jpg'))
# download the checkpoint from model zoo and put it in `checkpoints/`
checkpoint_file = 'work_dirs/solov2_tianchi_tuned.pth'
# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# # test a single image
# img = './WechatIMG14.jpeg'
# result, cost_time = inference_detector(model, img)
# show_result_ins(img, result, model.CLASSES, score_thr=0.25,
# out_file='./WechatIMG14_out.jpeg')
imgs = glob.glob('/workspace/tianchi/tianchiyusai/JPEGImages/621838/*.jpg')
# imgs = ['./test_imgs/14755.jpg']
# imgs = glob.glob('/home/versa/dataset/MSCOCO/aug_seg/val_imgs/*.*')
save_dir = './result'
if not os.path.exists(save_dir):
os.makedirs(save_dir)
total = 0
for i, img in enumerate(imgs):
name = img.split('/')[-1]
result, cost_time = inference_detector(model, img)
print(i, name, cost_time)
total += cost_time
try:
show_result_ins(img, result, model.CLASSES, score_thr=0.25,
out_file=os.path.join(save_dir, name))
except:
continue
print('average cost time: ', total / len(imgs))
import mmcv
import torch
# Decoupled solo
# config_file = '../configs/solo/decoupled_solo_r50_fpn_8gpu_3x.py'
# download the checkpoint from model zoo and put it in `checkpoints/`
# checkpoint_file = '../weights/DECOUPLED_SOLO_R50_3x.pth'
# Decoupled solo lite
# config_file = '../configs/solo/decoupled_solo_light_r50_fpn_8gpu_3x.py'
# download the checkpoint from model zoo and put it in `checkpoints/`
# checkpoint_file = '../weights/DECOUPLED_SOLO_LIGHT_R50_3x.pth'
# SOLOv2 lite
config_file = '../configs/solov2/solov2_light_448_r34_fpn_8gpu_3x.py'
# download the checkpoint from model zoo and put it in `checkpoints/`
checkpoint_file = '../weights/SOLOv2_LIGHT_448_R34_3x.pth'
# build the model from a config file and a checkpoint file
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# test a single image
img = 'images/demo.jpg'
result = inference_detector(model, img)
show_result_ins(img,
result,
model.CLASSES,
score_thr=0.25,
out_file="images/demo_out_torch.jpg")
