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')