def show_videos_with_trackbar(video_name, nets_dict, generate_img, flag_showIOU=False):
'''
:param video_name: list(["video path ", int])
:param nets_dict:
:param generate_img: bgr img -> tensor 等操作变换函数
:return:
'''
video = video_name[0]
cv2.namedWindow("video", cv2.NORM_HAMMING2)
cv2.resizeWindow("video", 1920, 1080)
cap = cv2.VideoCapture(video)
save_path = os.path.join("/home/xjx/data/videos/mask/mask_test_video/lta", video.split('/')[-1].split('.')[0])
frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 获得总帧数
# print cap.get(cv2.CAP_PROP_FPS) # 获得FPS
loop_flag = 0
pos = 0
cv2.createTrackbar('time', 'video', 0, frames, nothing)
trans_CropRoi = func.CropRoi(cfg)
while 1:
if loop_flag == pos: # 视频起始位置
loop_flag = loop_flag + 1
cv2.setTrackbarPos('time', 'video', loop_flag)
else: # 设置视频播放位置
pos = cv2.getTrackbarPos('time', 'video')
loop_flag = pos
cap.set(cv2.CAP_PROP_POS_FRAMES, pos) # 设置当前帧所在位置
ret, raw_img = cap.read()
# cv2.imshow("raw", raw_img)
img_tensor = generate_img(raw_img)
# im_show_all = dict(mix_img=[], mask_img=[], raw_img=[], person=[], back=[], change_back_addweight=[], change_back=[])# 组合图片
im_show_vis_type = list()
mask_list = list()
raw_list = list()
for net_name in nets_dict.keys():
net = nets_dict[net_name]
pre_dict = net(img_tensor)
im_show_vis_type.append(pre_dict[str(net.vis_type)]) # 每个网络选择自己的输出模式
# im_show_vis_type.append(pre_dict['mix_img']) # 每个网络选择自己的输出模式
# roiPred = trans_CropRoi(pre_dict, net, generate_img)
# im_show_vis_type.append(roiPred['mix_img'])
# mask_list.append(pre_dict['mask_img'][..., 0])
# mask_list.append(roiPred['mix_img'])
# raw_list.append(pre_dict['raw_img'
im_show_vis_type.append(raw_img) # 每个网络选择自己的输出模式
list_name = list(nets_dict.keys())
img_one = func.putImgsToOne(im_show_vis_type, list_name, 4, (0, 20), fonts=1,color=(0, 0, 255))
# mask_one = func.putImgsToOne(mask_list, list_name, 4, (0, 20), fonts=1,color=(0, 0, 255))
if flag_showIOU:
roi_img = func.vis_Image_from_iou(raw_img)
# roi_im_show_all = dict(mix_img=[], mask_img=[], raw_img=[], person=[], back=[], change_back_addweight=[], change_back=[])# 组合图片
roi_im_show_all = list()
if roi_img.shape[0] != 0: # 如果为空 shape为 (0,0,3)
roi_img_tensor = generate_img(roi_img)
for net_name in nets_dict.keys():
net = nets_dict[net_name]
pre_dict = nets_dict[net_name](roi_img_tensor)
roi_im_show_all.append(pre_dict[str(net.vis_type)]) # 每个网络选择自己的输出模式
roi_img_one = func.putImgsToOne(roi_im_show_all, list(nets_dict.keys()), 4, (0, 20), fonts=1,
color=(0, 0, 255))
cv2.imshow("Image", roi_img_one)
cv2.imshow("video", img_one)
# cv2.imshow("mask", mask_one)
key = cv2.waitKey(0)
if key == ord('q') or loop_flag == frames:
break
if key == ord('s'):
mmcv.imwrite(raw_list[-1], "%s/%06d.jpg" % (save_path, loop_flag))
mmcv.imwrite(mask_list[-1], "%s/%06d.png" % (save_path, loop_flag))
def main(cfg, nets_dict, flag):
'''
:param cfg:
:param nets_dict: {name: <class>Net_loader}
NetLoader __call__ return : dict(mask_img, mix_img, raw_img)
:return:
'''
generate_img = generate_img_test(cfg)
# imgs_dir = "/home/xjx/data/videos/mask/rvos/5_mask"
imgs_dir = "/home/lafe/work/mask/data/alldata/imgsave"
# imgs_list = ['49bbb9a205de862aa7dced16a6c22ddf_7.jpg', 'aa937c02e2119ec99dbac74ba0906ac3_3.jpg',
# '76a73c6827e2a643341d3a6e6a23c92b_3.jpg']
# imgs_list = [os.path.join(imgs_dir, i) for i in imgs_list]
# video_root = "/home/xjx/data/videos/mask/test_video"
# video_root = "/home/xjx/data/videos/mask/mask_test_video/PT_LB_Video"
video_root = "/home/xjx/data/videos/mask/mask_test_video/sz_dy"
# video_root = "/home/xjx/data/videos/mask/mask_test_video/PT_LB_Video" # 单人抖音视频
# video_root = "/home/xjx/data/videos/videoD/det_raw_video"
# videos = [[os.path.join(video_root, l), 0] for l in list(os.listdir(video_root))]
# videos = sorted(videos, key=lambda x: x)
# videos = [["/home/xjx/data/videos/mask/mask_test_video/sz_dy/PM_LB_JWJ_000056.mp4", 0]]
flag_showIOU = False
''' mix_img mask_img raw_img person back '''
# if flag_img:
if flag == 0:
if imgs_list is None:
imgs_list = list(os.path.join(imgs_dir, l) for l in list(os.listdir(imgs_dir)))
imgs_list = sorted(imgs_list , key=lambda x: x)
for img_dir in imgs_list:
try:
im = cv2.imread(img_dir)
try_test = im.shape
# im = im[478:955, 190:351, :]
except:
print("read error %s " % img_dir)
continue
img_tensor = generate_img(im) # 图片进行torch 格式转换
im_show_vis_type = list()
mask_list = []
for net_name in nets_dict.keys():
net = nets_dict[net_name]
pre_dict = net(img_tensor)
im_show_vis_type.append(pre_dict[str(net.vis_type)]) # 每个网络选择自己的输出模式
mask_list.append(pre_dict['mask_img'][...,0]) # 获得mask 单通道
img_one = func.putImgsToOne(im_show_vis_type, list(nets_dict.keys()), 4, (0, 20), fonts=1,
color=(0, 0, 255))
cv2.imshow('img', img_one)
k = cv2.waitKey(0)
if k == ord('q'):
cv2.destroyAllWindows()
break
if k == ord('s'):
cv2.imwrite('/home/xjx/data/videos/mask/rvos/5_mask/0.png', mask_list[0])
print('Test finished')
# if flag_video:
if flag == 1:
for video_name in videos:
try:
show_videos_with_trackbar(video_name, nets_dict, generate_img, flag_showIOU)
except:
print("%s finised!" %video_name[0])
cv2.destroyAllWindows()
# cap
if flag == 2:
try:
cap = cv2.VideoCapture(0)
ret, im = cap.read()
zzz = im.shape
except:
cap = cv2.VideoCapture(1)
cap_frame = 0
cv2.namedWindow("img", cv2.NORM_HAMMING2)
while 1:
ret, im = cap.read()
# 1. ------------------增广操作------------------------------
# cv2.imshow('raw', im)
img_tensor = generate_img(im)
im_show_vis_type = list()
for net_name in nets_dict.keys():
net = nets_dict[net_name]
pre_dict = net(img_tensor)
im_show_vis_type.append(pre_dict[str(net.vis_type)]) # 每个网络选择自己的输出模式
img_one = func.putImgsToOne(im_show_vis_type, list(nets_dict.keys()), 4, (0, 20), fonts=1,
color=(0, 0, 255))
# --------------- 马赛克 --------------------
# ii = im_show_all['person'][1]
# h, w, c = ii.shape
# im_show_all['person'][1] = do_mosaic(ii, 0, 0, w, h, 20)
# cv2.resizeWindow("img", 1920, 1080)
cv2.imshow("img", img_one)
key = cv2.waitKey(5)
if key == ord('q'):
cv2.destroyAllWindows()
break
def inference(cfg, nets_dict, imgs_dir=None, mask_save=None, flag_save=False):
'''
:param cfg:
:param nets_dict: {name: <class>Net_loader}
NetLoader __call__ return : dict(mask_img, mix_img, raw_img)
:return:
'''
if flag_save:
os.makedirs(mask_save, exist_ok=True)
generate_img = generate_img_test(cfg)
flag_showIOU = False
''' mix_img mask_img raw_img person back '''
# if flag_img:
imgs_list = list(
os.path.join(imgs_dir, l) for l in list(os.listdir(imgs_dir)))
imgs_list = sorted(imgs_list, key=lambda x: x)
vis = []
for img_dir in tqdm(imgs_list[0:20]):
try:
im = cv2.imread(img_dir)
basename = os.path.basename(img_dir)
try_test = im.shape
# im = im[478:955, 190:351, :]
except:
print("read error %s " % img_dir)
continue
img_tensor = generate_img(im) # 图片进行torch 格式转换
im_show_vis_type = list()
mask_list = []
for net_name in nets_dict.keys():
net = nets_dict[net_name]
s = time.time()
pre_dict = net(img_tensor)
print(' * infere time: ', time.time() - s)
# im_show_vis_type.append(pre_dict[str(net.vis_type)]) # 每个网络选择自己的输出模式
for k, v in pre_dict.items():
im_show_vis_type.append(v)
mask_list.append(pre_dict['mask_img']) # 获得mask 单通道
if flag_save:
savename = os.path.join(mask_save, basename.replace('jpg', 'png'))
mask = mask_list[-1]
if len(mask.shape) == 3:
mask = mask[..., 0]
cv2.imwrite(savename, mask)
else:
img_one = func.putImgsToOne(im_show_vis_type,
list(nets_dict.keys()),
4, (0, 20),
fonts=1,
color=(0, 0, 255))
vis.append(img_one)
# cv2.imshow('img', img_one)
# k = cv2.waitKey(0)
# if k == ord('q'):
# cv2.destroyAllWindows()
# break
if not flag_save:
loopshow(vis)
print('Test finished')