''' 只crop可扣除一只手的图，没有resize。 example: python utils/crop_resize_no_another_hand.py --video_id 20230818_04_old TODO：检查为什么运行到后面会越来越慢。似乎是卡在了load_bg_img，可以测一下这一行的时间。 re:好像只是单纯有点慢 ''' import os import sys sys.path.append('.') import cv2 import numpy as np import os.path as osp from tqdm import tqdm import pickle import multiprocessing as mlp from utils.hoi_io2 import get_downsampled_seg_infos_batch, get_seg_infos_batch3, load_bg_img, get_downsampled_seg_infos_batch_v2, read_init_crop, cal_represent_frame_list, get_downsampled_seg_infos_batch_v2_acc_batch from utils.scandir import scandir import argparse from time import time import json def crop_from_mask(video_id: str, camera_list: str, frame_list: list[str]): ''' 当前帧如果没有某个mask，则用上一帧的bbox进行crop。 ''' seg, downsample_factor = get_downsampled_seg_infos_batch(video_id, frame_list, camera_list) right_hand_seg = np.where(seg == 1, 1, 0).astype(np.uint8) left_hand_seg = np.where(seg == 2, 1, 0).astype(np.uint8) last_right_min_x = None last_right_max_x = None last_right_min_y = None last_right_max_y = None last_right_meanh = None last_right_meanw = None last_left_min_x = None last_left_max_x = None last_left_min_y = None last_left_max_y = None last_left_meanh = None last_left_meanw = None MAX_HEIGHT = 4095 MAX_WIDTH = 2999 MARGIN_SIZE = 50 # TODO：以下的x和y定义反了，需要更改 for camera_idx, camera_id in enumerate(camera_list): # left_hand_crop_dir = osp.join(root, video_id, 'crop_imgs_left_hand', camera_id) # right_hand_crop_dir = osp.join(root, video_id, 'crop_imgs_right_hand', camera_id) left_hand_crop_dir = osp.join(root, video_id, 'crop', 'left_hand', camera_id) right_hand_crop_dir = osp.join(root, video_id, 'crop', 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) for frame_idx, frame_id in tqdm(enumerate(frame_list)): mask_right_hand = right_hand_seg[frame_idx, camera_idx] mask_left_hand = left_hand_seg[frame_idx, camera_idx] # upsample mask_right_hand = cv2.resize(mask_right_hand, (4096, 3000), interpolation=cv2.INTER_LINEAR) mask_left_hand = cv2.resize(mask_left_hand, (4096, 3000), interpolation=cv2.INTER_LINEAR) # 读入图片 # timer1 = time() img = load_bg_img(video_id, camera_id, frame_id) # timer2 = time() # print(timer2 - timer1) # 以下为右手 mask_right_hand_idx = np.nonzero(mask_right_hand) if len(mask_right_hand_idx[0]) != 0: right_minh = np.min(mask_right_hand_idx[0]) right_maxh = np.max(mask_right_hand_idx[0]) right_minw = np.min(mask_right_hand_idx[1]) right_maxw = np.max(mask_right_hand_idx[1]) right_midh = (right_maxh + right_minh) // 2 right_midw = (right_maxw + right_minw) // 2 right_meanh = np.mean(mask_right_hand_idx[0]) right_meanw = np.mean(mask_right_hand_idx[1]) right_length = max(right_maxh - right_minh, right_maxw - right_minw) right_min_x = max(0, right_midh - right_length // 2 - MARGIN_SIZE) right_max_x = min(MAX_WIDTH, right_midh + right_length // 2 + MARGIN_SIZE) right_min_y = max(0, right_midw - right_length // 2 - MARGIN_SIZE) right_max_y = min(MAX_HEIGHT, right_midw + right_length // 2 + MARGIN_SIZE) last_right_min_x = right_min_x last_right_max_x = right_max_x last_right_min_y = right_min_y last_right_max_y = right_max_y last_right_meanh = right_meanh last_right_meanw = right_meanw elif frame_idx != 0 and len(mask_right_hand_idx[0]) == 0: right_min_x = last_right_min_x right_max_x = last_right_max_x right_min_y = last_right_min_y right_max_y = last_right_max_y right_meanh = last_right_meanh right_meanw = last_right_meanw # else: # print('no hand mask in the first frame!') # exit(1) else: print('no hand mask in the first frame of this sub_frame_list!') print('right_hand', camera_id, frame_id) if frame_idx == 0 and frame_id != '00001': info_path = osp.join(right_hand_crop_dir, camera_id + '_' + str(int(frame_id)-1).zfill(5) + '_crop_info.pkl') assert os.path.exists(info_path) with open(info_path, 'rb') as f: right_min_x, right_max_x, right_min_y, right_max_y, right_meanh, right_meanw = pickle.load(f) last_right_min_x = right_min_x last_right_max_x = right_max_x last_right_min_y = right_min_y last_right_max_y = right_max_y last_right_meanh = right_meanh last_right_meanw = right_meanw else: print('GG, no hand mask in the first frame!') exit(1) # 将左手变成白色 right_crop = img.copy() right_crop[mask_left_hand == 1] = 255.0 # 居中 right_crop = right_crop[right_min_x: right_max_x, right_min_y: right_max_y] right_crop_info = [right_min_x, right_max_x, right_min_y, right_max_y, right_meanh, right_meanw] with open(osp.join(right_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(right_crop_info, f) cv2.imwrite(osp.join(right_hand_crop_dir, camera_id + '_' +frame_id + '.png'), right_crop) # 以下为左手 mask_left_hand_idx = np.nonzero(mask_left_hand) # test_img = np.zeros((3000, 4096, 3)) # for y, x in zip(mask_left_hand_idx[0], mask_left_hand_idx[1]): # cv2.circle(test_img, (x, y), 5, (255, 0, 0), -1) # cv2.imwrite('./test.png',test_img) if len(mask_left_hand_idx[0]) != 0: left_minh = np.min(mask_left_hand_idx[0]) left_maxh = np.max(mask_left_hand_idx[0]) left_minw = np.min(mask_left_hand_idx[1]) left_maxw = np.max(mask_left_hand_idx[1]) left_midh = (left_maxh + left_minh) // 2 left_midw = (left_maxw + left_minw) // 2 left_meanh = np.mean(mask_left_hand_idx[0]) left_meanw = np.mean(mask_left_hand_idx[1]) left_length = max(left_maxh - left_minh, left_maxw - left_minw) left_min_x = max(0, left_midh - left_length // 2 - MARGIN_SIZE) left_max_x = min(MAX_WIDTH, left_midh + left_length // 2 + MARGIN_SIZE) left_min_y = max(0, left_midw - left_length // 2 - MARGIN_SIZE) left_max_y = min(MAX_HEIGHT, left_midw + left_length // 2 + MARGIN_SIZE) last_left_min_x = left_min_x last_left_max_x = left_max_x last_left_min_y = left_min_y last_left_max_y = left_max_y last_left_meanh = left_meanh last_left_meanw = left_meanw elif frame_idx != 0 and len(mask_left_hand_idx[0]) == 0: left_min_x = last_left_min_x left_max_x = last_left_max_x left_min_y = last_left_min_y left_max_y = last_left_max_y left_meanh = last_left_meanh left_meanw = last_left_meanw else: print('no hand mask in the first frame of this sub_frame_list!') print('left_hand', camera_id, frame_id) if frame_idx == 0 and frame_id != '00001': info_path = osp.join(left_hand_crop_dir, camera_id + '_' + str(int(frame_id)-1).zfill(5) + '_crop_info.pkl') assert os.path.exists(info_path) with open(info_path, 'rb') as f: left_min_x, left_max_x, left_min_y, left_max_y, left_meanh, left_meanw = pickle.load(f) last_left_min_x = left_min_x last_left_max_x = left_max_x last_left_min_y = left_min_y last_left_max_y = left_max_y last_left_meanh = left_meanh last_left_meanw = left_meanw else: print(frame_idx, frame_id) print('GG, no hand mask in the first frame!') exit(1) # 将右手变成白色 left_crop = img.copy() left_crop[mask_right_hand == 1] = 255.0 # 居中 left_crop = left_crop[left_min_x: left_max_x, left_min_y: left_max_y] left_crop_info = [left_min_x, left_max_x, left_min_y, left_max_y, left_meanh, left_meanw] with open(osp.join(left_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(left_crop_info, f) cv2.imwrite(osp.join(left_hand_crop_dir, camera_id + '_' + frame_id + '.png'), left_crop) def crop(img, min_h, max_h, min_w, max_w, black_min_h=None, black_max_h=None, black_min_w=None, black_max_w=None): img = img.copy() mid_h = (min_h + max_h) // 2 mid_w = (min_w + max_w) // 2 length = max(max_h - min_h, max_w - min_w) if [black_min_h, black_max_h, black_min_w, black_max_w] == [not None, not None, not None, not None]: img[black_min_h: black_max_h, black_min_w, black_max_w] = 255.0 crop = img[min_h: max_h, min_w: max_w] return crop def apply_factor_and_margin(min_h, max_h, min_w, max_w, factor, margin_size, MAX_H, MAX_W): # # avoid overflow # min_h_ = min_h.astype(np.uint16) # max_h_ = max_h.astype(np.uint16) # min_w_ = min_w.astype(np.uint16) # max_w_ = max_w.astype(np.uint16) # mid_h = ((min_h_ + max_h_) // 2).astype(np.uint8) # mid_w = ((min_w_ + max_w_) // 2).astype(np.uint8) mid_h = (min_h + max_h) // 2 mid_w = (min_w + max_w) // 2 length = max(max_h - min_h, max_w - min_w) length = int(length * factor) half_length = length // 2 min_h = max(0, mid_h - half_length - margin_size, 0) max_h = min(MAX_H, mid_h + half_length + margin_size) min_w = max(0, mid_w - half_length - margin_size, 0) max_w = min(MAX_W, mid_w + half_length + margin_size) return min_h, max_h, min_w, max_w, mid_h, mid_w def crop_hand_from_pos(root, video_id: str, camera_id: str, frame_id: str, pos, crop_factor = 1, margin_size = 0): H = 3000 W = 4096 MAX_HEIGHT = H - 1 MAX_WIDTH = W - 1 assert pos.shape == (2, 4) # assert np.all(pos[...] >= 0), print(pos) # assert np.all(pos[..., [0, 1]] <= MAX_HEIGHT), print(pos) # assert np.all(pos[..., [2, 3]] <= MAX_WIDTH), print(pos) if np.all(pos[...] >= 0) and np.all(pos[..., [0, 1]] <= MAX_HEIGHT) and np.all(pos[..., [2, 3]] <= MAX_WIDTH): left_hand_crop_dir = osp.join(root, video_id, 'crop', 'left_hand', camera_id) right_hand_crop_dir = osp.join(root, video_id, 'crop', 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) bg = load_bg_img(video_id, camera_id, frame_id) # crop info right_crop_info_path = osp.join(right_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl') (right_min_h, right_max_h, right_min_w, right_max_w) = pos[0] right_min_h, right_max_h, right_min_w, right_max_w, right_mid_h, right_mid_w = apply_factor_and_margin(right_min_h, right_max_h, right_min_w, right_max_w, crop_factor, margin_size, MAX_HEIGHT, MAX_WIDTH) right_crop_info = [right_min_h, right_max_h, right_min_w, right_max_w, right_mid_h, right_mid_w] with open(right_crop_info_path, 'wb') as f: pickle.dump(right_crop_info, f) left_crop_info_path = osp.join(left_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl') (left_min_h, left_max_h, left_min_w, left_max_w) = pos[1] left_min_h, left_max_h, left_min_w, left_max_w, left_mid_h, left_mid_w = apply_factor_and_margin(left_min_h, left_max_h, left_min_w, left_max_w, crop_factor, margin_size, MAX_HEIGHT, MAX_WIDTH) left_crop_info = [left_min_h, left_max_h, left_min_w, left_max_w, left_mid_h, left_mid_w] with open(left_crop_info_path, 'wb') as f: pickle.dump(left_crop_info, f) right_crop = crop(bg, right_min_h, right_max_h, right_min_w, right_max_w, left_min_h, left_max_h, left_min_w, left_max_w) right_crop_path = osp.join(right_hand_crop_dir, camera_id + '_' + frame_id + '.png') cv2.imwrite(right_crop_path, right_crop) # left hand left_crop = crop(bg, left_min_h, left_max_h, left_min_w, left_max_w, right_min_h, right_max_h, right_min_w, right_max_w) left_crop_path = osp.join(left_hand_crop_dir, camera_id + '_' + frame_id + '.png') cv2.imwrite(left_crop_path, left_crop) def crop_hand_from_pos_acc(root, local_root, date: str, crop_save_exp_name, crop_info_save_exp_name, video_id: str, camera_list: str, frame_list: str, represent_frame_id, img_batch, pos_batch, crop_factor = 1, margin_size = 0): H = 3000 W = 4096 MAX_HEIGHT = H - 1 MAX_WIDTH = W - 1 # assert pos.shape == (2, 4) # assert np.all(pos[...] >= 0), print(pos) # assert np.all(pos[..., [0, 1]] <= MAX_HEIGHT), print(pos) # assert np.all(pos[..., [2, 3]] <= MAX_WIDTH), print(pos) # TODO right_crop_info_batch = {} left_crop_info_batch = {} for f_idx, frame_id in enumerate(frame_list): right_crop_info_batch[frame_id] = {} left_crop_info_batch[frame_id] = {} for c_idx, camera_id in enumerate(camera_list): if camera_id not in pos_batch[frame_id].keys(): continue pos = pos_batch[frame_id][camera_id] if np.all(pos[...] >= 0) and np.all(pos[..., [0, 1]] <= MAX_HEIGHT) and np.all(pos[..., [2, 3]] <= MAX_WIDTH): left_hand_crop_dir = osp.join(local_root, date, video_id, crop_save_exp_name, 'left_hand', camera_id) right_hand_crop_dir = osp.join(local_root, date, video_id, crop_save_exp_name, 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) bg = img_batch[f_idx][c_idx].copy() # crop info (right_min_h, right_max_h, right_min_w, right_max_w) = pos[0] right_min_h, right_max_h, right_min_w, right_max_w, right_mid_h, right_mid_w = apply_factor_and_margin(right_min_h, right_max_h, right_min_w, right_max_w, crop_factor, margin_size, MAX_HEIGHT, MAX_WIDTH) right_crop_info = [right_min_h, right_max_h, right_min_w, right_max_w, right_mid_h, right_mid_w] right_crop_info_batch[frame_id][camera_id] = right_crop_info # right_crop_info_path = osp.join(right_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl') # with open(right_crop_info_path, 'wb') as f: # pickle.dump(right_crop_info, f) (left_min_h, left_max_h, left_min_w, left_max_w) = pos[1] left_min_h, left_max_h, left_min_w, left_max_w, left_mid_h, left_mid_w = apply_factor_and_margin(left_min_h, left_max_h, left_min_w, left_max_w, crop_factor, margin_size, MAX_HEIGHT, MAX_WIDTH) left_crop_info = [left_min_h, left_max_h, left_min_w, left_max_w, left_mid_h, left_mid_w] left_crop_info_batch[frame_id][camera_id] = left_crop_info # left_crop_info_path = osp.join(left_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl') # with open(left_crop_info_path, 'wb') as f: # pickle.dump(left_crop_info, f) right_crop = crop(bg, right_min_h, right_max_h, right_min_w, right_max_w, left_min_h, left_max_h, left_min_w, left_max_w) right_crop_path = osp.join(right_hand_crop_dir, camera_id + '_' + frame_id + '.png') cv2.imwrite(right_crop_path, right_crop) # left hand left_crop = crop(bg, left_min_h, left_max_h, left_min_w, left_max_w, right_min_h, right_max_h, right_min_w, right_max_w) left_crop_path = osp.join(left_hand_crop_dir, camera_id + '_' + frame_id + '.png') cv2.imwrite(left_crop_path, left_crop) right_crop_info_dir = osp.join(local_root, date, video_id, crop_info_save_exp_name, 'right_hand') os.makedirs(right_crop_info_dir, exist_ok=True) left_crop_info_dir = osp.join(local_root, date, video_id, crop_info_save_exp_name, 'left_hand') os.makedirs(left_crop_info_dir, exist_ok=True) right_crop_info_path = osp.join(right_crop_info_dir, represent_frame_id + '_crop_info.pkl') with open(right_crop_info_path, 'wb') as f: pickle.dump(right_crop_info_batch, f) left_crop_info_path = osp.join(left_crop_info_dir, represent_frame_id + '_crop_info.pkl') with open(left_crop_info_path, 'wb') as f: pickle.dump(left_crop_info_batch, f) def crop_from_init(root, date, video_id: str, camera_list: str, crop_factor = 1, margin_size = 0): ''' 手动标注第一二帧，从src中读取标注数据。 crop一只手时，将另一只手的bbox所在区域全部置为黑。 TODO: pos的shape有问题，要改 ''' frame_list = [str(i).zfill(5) for i in range(1, 3)] for camera_idx, camera_id in enumerate(camera_list): left_hand_crop_dir = osp.join(root, date, video_id, 'crop', 'left_hand', camera_id) right_hand_crop_dir = osp.join(root, date, video_id, 'crop', 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) for frame_idx, frame_id in enumerate(frame_list): pos = read_init_crop(video_id, camera_id, frame_id) crop_hand_from_pos(video_id, camera_id, frame_id, pos, crop_factor, margin_size) def crop_from_mask_v2(root, video_id: str, from_exp_name, camera_list: str, frame_list: list[str], rgb_batch, crop_factor = 1, margin_size = 50): # for camera_id in camera_list: # left_hand_crop_dir = os.path.join(root, video_id, 'crop', 'left_hand', camera_id) # right_hand_crop_dir = os.path.join(root, video_id, 'crop', 'right_hand', camera_id) # os.makedirs(left_hand_crop_dir, exist_ok=True) # os.makedirs(right_hand_crop_dir, exist_ok=True) seg, downsample_factor = get_downsampled_seg_infos_batch_v2(video_id, from_exp_name, frame_list, camera_list) right_hand_seg = np.where(seg == 1, 1, 0).astype(np.uint8) left_hand_seg = np.where(seg == 2, 1, 0).astype(np.uint8) right_hand_valid_camera_dict = {} left_hand_valid_camera_dict = {} for frame in frame_list: right_hand_valid_camera_dict[frame] = [] left_hand_valid_camera_dict[frame] = [] for camera_idx, camera_id in enumerate(camera_list): left_hand_crop_dir = os.path.join(root, video_id, 'crop', 'left_hand', camera_id) right_hand_crop_dir = os.path.join(root, video_id, 'crop', 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) for frame_idx, frame_id in tqdm(enumerate(frame_list)): mask_right_hand = right_hand_seg[frame_idx, camera_idx] mask_left_hand = left_hand_seg[frame_idx, camera_idx] # 读入图片 img = load_bg_img(video_id, camera_id, frame_id) # H = 3000 # W = 4096 (H, W) = img.shape[:2] MAX_WIDTH = W - 1 MAX_HEIGHT = H - 1 # upsample mask mask_right_hand = cv2.resize(mask_right_hand, (W, H), interpolation=cv2.INTER_NEAREST) mask_left_hand = cv2.resize(mask_left_hand, (W, H), interpolation=cv2.INTER_NEAREST) # 以下为右手 mask_right_hand_idx = np.nonzero(mask_right_hand) if len(mask_right_hand_idx[0]) != 0: right_hand_valid_camera_dict[frame_id].append(camera_id) right_minh = np.min(mask_right_hand_idx[0]) right_maxh = np.max(mask_right_hand_idx[0]) right_minw = np.min(mask_right_hand_idx[1]) right_maxw = np.max(mask_right_hand_idx[1]) right_midh = (right_maxh + right_minh) // 2 right_midw = (right_maxw + right_minw) // 2 right_meanh = np.mean(mask_right_hand_idx[0]) right_meanw = np.mean(mask_right_hand_idx[1]) right_length = int(max(right_maxh - right_minh, right_maxw - right_minw) * crop_factor) half_length = right_length // 2 right_min_y = max(0, right_midh - half_length - margin_size) right_max_y = min(MAX_HEIGHT, right_midh + half_length + margin_size) right_min_x = max(0, right_midw - half_length - margin_size) right_max_x = min(MAX_WIDTH, right_midw + half_length + margin_size) # 将左手变成白色 right_crop = img.copy() right_crop[mask_left_hand == 1] = 255.0 # 居中 right_crop = right_crop[right_min_y: right_max_y, right_min_x: right_max_x] right_crop_info = [right_min_y, right_max_y, right_min_x, right_max_x, right_meanh, right_meanw] with open(os.path.join(right_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(right_crop_info, f) cv2.imwrite(os.path.join(right_hand_crop_dir, camera_id + '_' +frame_id + '.png'), right_crop) # 以下为左手 mask_left_hand_idx = np.nonzero(mask_left_hand) if len(mask_left_hand_idx[0]) != 0: left_hand_valid_camera_dict[frame_id].append(camera_id) left_minh = np.min(mask_left_hand_idx[0]) left_maxh = np.max(mask_left_hand_idx[0]) left_minw = np.min(mask_left_hand_idx[1]) left_maxw = np.max(mask_left_hand_idx[1]) left_midh = (left_maxh + left_minh) // 2 left_midw = (left_maxw + left_minw) // 2 left_meanh = np.mean(mask_left_hand_idx[0]) left_meanw = np.mean(mask_left_hand_idx[1]) left_length = max(left_maxh - left_minh, left_maxw - left_minw) half_length = left_length // 2 left_min_y = max(0, left_midh - half_length - margin_size) left_max_y = min(MAX_HEIGHT, left_midh + half_length + margin_size) left_min_x = max(0, left_midw - half_length - margin_size) left_max_x = min(MAX_WIDTH, left_midw + half_length + margin_size) # 将右手变成白色 left_crop = img.copy() left_crop[mask_right_hand == 1] = 255.0 # 居中 left_crop = left_crop[left_min_y: left_max_y, left_min_x: left_max_x] left_crop_info = [left_min_y, left_max_y, left_min_x, left_max_x, left_meanh, left_meanw] with open(os.path.join(left_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(left_crop_info, f) cv2.imwrite(os.path.join(left_hand_crop_dir, camera_id + '_' + frame_id + '.png'), left_crop) # save valid and invalid info right_hand_invalid_camera_dict = {} left_hand_invalid_camera_dict = {} for frame in frame_list: right_hand_invalid_camera_dict[frame] = [camera for camera in camera_list if camera not in right_hand_valid_camera_dict[frame]] left_hand_invalid_camera_dict[frame] = [camera for camera in camera_list if camera not in left_hand_valid_camera_dict[frame]] log_path = os.path.join(root, video_id, 'crop', 'log.json') if os.path.exists(log_path): with open(log_path, 'r') as f: log = json.load(f) log['right_hand_valid_camera_dict'].update(right_hand_valid_camera_dict) log['left_hand_valid_camera_dict'].update(left_hand_valid_camera_dict) log['right_hand_invalid_camera_dict'].update(right_hand_invalid_camera_dict) log['left_hand_invalid_camera_dict'].update(left_hand_invalid_camera_dict) else: log = {} log['right_hand_valid_camera_dict'] = right_hand_valid_camera_dict log['left_hand_valid_camera_dict'] = left_hand_valid_camera_dict log['right_hand_invalid_camera_dict'] = right_hand_invalid_camera_dict log['left_hand_invalid_camera_dict'] = left_hand_invalid_camera_dict with open(log_path, 'w') as f: json.dump(log, f) return right_hand_valid_camera_dict, left_hand_valid_camera_dict, right_hand_invalid_camera_dict, left_hand_invalid_camera_dict def crop_from_mask_v3(root, video_id: str, from_exp_name, camera_list: str, frame_list: list[str], rgb_batch, crop_factor = 1, margin_size = 50): ''' 增加了rgb_batch，可以不用从本地再读图片。 ''' # for camera_id in camera_list: # left_hand_crop_dir = os.path.join(root, video_id, 'crop', 'left_hand', camera_id) # right_hand_crop_dir = os.path.join(root, video_id, 'crop', 'right_hand', camera_id) # os.makedirs(left_hand_crop_dir, exist_ok=True) # os.makedirs(right_hand_crop_dir, exist_ok=True) seg, downsample_factor = get_downsampled_seg_infos_batch_v2(video_id, from_exp_name, frame_list, camera_list) right_hand_seg = np.where(seg == 1, 1, 0).astype(np.uint8) left_hand_seg = np.where(seg == 2, 1, 0).astype(np.uint8) right_hand_valid_camera_dict = {} left_hand_valid_camera_dict = {} for frame in frame_list: right_hand_valid_camera_dict[frame] = [] left_hand_valid_camera_dict[frame] = [] for camera_idx, camera_id in enumerate(camera_list): left_hand_crop_dir = os.path.join(root, video_id, 'crop', 'left_hand', camera_id) right_hand_crop_dir = os.path.join(root, video_id, 'crop', 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) for frame_idx, frame_id in tqdm(enumerate(frame_list)): mask_right_hand = right_hand_seg[frame_idx, camera_idx] mask_left_hand = left_hand_seg[frame_idx, camera_idx] # 读入图片 # img = rgb_batch[int(frame_id)-1,camera_idx].copy() # img = rgb_batch[int(frame_id)-1][camera_idx].copy() img = rgb_batch[frame_idx][camera_idx].copy() # img = load_bg_img(video_id, camera_id, frame_id) # H = 3000 # W = 4096 (H, W) = img.shape[:2] MAX_WIDTH = W - 1 MAX_HEIGHT = H - 1 # upsample mask mask_right_hand = cv2.resize(mask_right_hand, (W, H), interpolation=cv2.INTER_NEAREST) mask_left_hand = cv2.resize(mask_left_hand, (W, H), interpolation=cv2.INTER_NEAREST) # 以下为右手 mask_right_hand_idx = np.nonzero(mask_right_hand) if len(mask_right_hand_idx[0]) != 0: right_hand_valid_camera_dict[frame_id].append(camera_id) right_minh = np.min(mask_right_hand_idx[0]) right_maxh = np.max(mask_right_hand_idx[0]) right_minw = np.min(mask_right_hand_idx[1]) right_maxw = np.max(mask_right_hand_idx[1]) right_midh = (right_maxh + right_minh) // 2 right_midw = (right_maxw + right_minw) // 2 right_meanh = np.mean(mask_right_hand_idx[0]) right_meanw = np.mean(mask_right_hand_idx[1]) right_length = int(max(right_maxh - right_minh, right_maxw - right_minw) * crop_factor) half_length = right_length // 2 right_min_y = max(0, right_midh - half_length - margin_size) right_max_y = min(MAX_HEIGHT, right_midh + half_length + margin_size) right_min_x = max(0, right_midw - half_length - margin_size) right_max_x = min(MAX_WIDTH, right_midw + half_length + margin_size) # 将左手变成白色 right_crop = img.copy() right_crop[mask_left_hand == 1] = 255.0 # 居中 right_crop = right_crop[right_min_y: right_max_y, right_min_x: right_max_x] right_crop_info = [right_min_y, right_max_y, right_min_x, right_max_x, right_meanh, right_meanw] with open(os.path.join(right_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(right_crop_info, f) cv2.imwrite(os.path.join(right_hand_crop_dir, camera_id + '_' +frame_id + '.png'), right_crop) # 以下为左手 mask_left_hand_idx = np.nonzero(mask_left_hand) if len(mask_left_hand_idx[0]) != 0: left_hand_valid_camera_dict[frame_id].append(camera_id) left_minh = np.min(mask_left_hand_idx[0]) left_maxh = np.max(mask_left_hand_idx[0]) left_minw = np.min(mask_left_hand_idx[1]) left_maxw = np.max(mask_left_hand_idx[1]) left_midh = (left_maxh + left_minh) // 2 left_midw = (left_maxw + left_minw) // 2 left_meanh = np.mean(mask_left_hand_idx[0]) left_meanw = np.mean(mask_left_hand_idx[1]) left_length = max(left_maxh - left_minh, left_maxw - left_minw) half_length = left_length // 2 left_min_y = max(0, left_midh - half_length - margin_size) left_max_y = min(MAX_HEIGHT, left_midh + half_length + margin_size) left_min_x = max(0, left_midw - half_length - margin_size) left_max_x = min(MAX_WIDTH, left_midw + half_length + margin_size) # 将右手变成白色 left_crop = img.copy() left_crop[mask_right_hand == 1] = 255.0 # 居中 left_crop = left_crop[left_min_y: left_max_y, left_min_x: left_max_x] left_crop_info = [left_min_y, left_max_y, left_min_x, left_max_x, left_meanh, left_meanw] with open(os.path.join(left_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(left_crop_info, f) cv2.imwrite(os.path.join(left_hand_crop_dir, camera_id + '_' + frame_id + '.png'), left_crop) # save valid and invalid info right_hand_invalid_camera_dict = {} left_hand_invalid_camera_dict = {} for frame in frame_list: right_hand_invalid_camera_dict[frame] = [camera for camera in camera_list if camera not in right_hand_valid_camera_dict[frame]] left_hand_invalid_camera_dict[frame] = [camera for camera in camera_list if camera not in left_hand_valid_camera_dict[frame]] log_path = os.path.join(root, video_id, 'crop', 'log.json') if os.path.exists(log_path): with open(log_path, 'r') as f: log = json.load(f) log['right_hand_valid_camera_dict'].update(right_hand_valid_camera_dict) log['left_hand_valid_camera_dict'].update(left_hand_valid_camera_dict) log['right_hand_invalid_camera_dict'].update(right_hand_invalid_camera_dict) log['left_hand_invalid_camera_dict'].update(left_hand_invalid_camera_dict) else: log = {} log['right_hand_valid_camera_dict'] = right_hand_valid_camera_dict log['left_hand_valid_camera_dict'] = left_hand_valid_camera_dict log['right_hand_invalid_camera_dict'] = right_hand_invalid_camera_dict log['left_hand_invalid_camera_dict'] = left_hand_invalid_camera_dict with open(log_path, 'w') as f: json.dump(log, f) return right_hand_valid_camera_dict, left_hand_valid_camera_dict, right_hand_invalid_camera_dict, left_hand_invalid_camera_dict def crop_from_mask_v3_acc_batch(root, local_root, date, video_id: str, from_exp_name, camera_list: str, represent_frame_id, frame_list: list[str], BATCH_SIZE, rgb_batch, crop_factor = 1, margin_size = 50): ''' 增加了rgb_batch，可以不用从本地再读图片。 每一个batch的文件分左右手写在一起。 ''' seg, downsample_factor = get_downsampled_seg_infos_batch_v2_acc_batch(root, date, video_id, from_exp_name, frame_list, camera_list, represent_frame_id=represent_frame_id) right_hand_seg = np.where(seg == 1, 1, 0).astype(np.uint8) left_hand_seg = np.where(seg == 2, 1, 0).astype(np.uint8) right_hand_valid_camera_dict = {} left_hand_valid_camera_dict = {} for frame in frame_list: right_hand_valid_camera_dict[frame] = [] left_hand_valid_camera_dict[frame] = [] # represent_relation = cal_represent_frame_list(BATCH_SIZE, frame_list) # represent_keys = list(represent_relation.keys()) # assert len(represent_keys) == 1 # represent_frame_id = represent_keys[0] left_crop_info_batch = {} left_crop_batch = {} right_crop_info_batch = {} right_crop_batch = {} for camera_idx, camera_id in enumerate(camera_list): left_hand_crop_dir = os.path.join(local_root, date, video_id, 'crop', 'left_hand', camera_id) right_hand_crop_dir = os.path.join(local_root, date, video_id, 'crop', 'right_hand', camera_id) os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) for frame_idx, frame_id in tqdm(enumerate(frame_list)): mask_right_hand = right_hand_seg[frame_idx, camera_idx] mask_left_hand = left_hand_seg[frame_idx, camera_idx] # 读入图片 img = rgb_batch[frame_idx][camera_idx].copy() # H = 3000 # W = 4096 (H, W) = img.shape[:2] MAX_WIDTH = W - 1 MAX_HEIGHT = H - 1 # upsample mask mask_right_hand = cv2.resize(mask_right_hand, (W, H), interpolation=cv2.INTER_NEAREST) mask_left_hand = cv2.resize(mask_left_hand, (W, H), interpolation=cv2.INTER_NEAREST) # 以下为右手 mask_right_hand_idx = np.nonzero(mask_right_hand) if len(mask_right_hand_idx[0]) != 0: right_hand_valid_camera_dict[frame_id].append(camera_id) right_minh = np.min(mask_right_hand_idx[0]) right_maxh = np.max(mask_right_hand_idx[0]) right_minw = np.min(mask_right_hand_idx[1]) right_maxw = np.max(mask_right_hand_idx[1]) right_midh = (right_maxh + right_minh) // 2 right_midw = (right_maxw + right_minw) // 2 right_meanh = np.mean(mask_right_hand_idx[0]) right_meanw = np.mean(mask_right_hand_idx[1]) right_length = int(max(right_maxh - right_minh, right_maxw - right_minw) * crop_factor) half_length = right_length // 2 right_min_y = max(0, right_midh - half_length - margin_size) right_max_y = min(MAX_HEIGHT, right_midh + half_length + margin_size) right_min_x = max(0, right_midw - half_length - margin_size) right_max_x = min(MAX_WIDTH, right_midw + half_length + margin_size) # 将左手变成白色 right_crop = img.copy() right_crop[mask_left_hand == 1] = 255.0 # 居中 right_crop = right_crop[right_min_y: right_max_y, right_min_x: right_max_x] right_crop_info = [right_min_y, right_max_y, right_min_x, right_max_x, right_meanh, right_meanw] if frame_id not in right_crop_info_batch: right_crop_info_batch[frame_id] = {} right_crop_info_batch[frame_id][camera_id] = right_crop_info if frame_id not in right_crop_batch: right_crop_batch[frame_id] = {} right_crop_batch[frame_id][camera_id] = right_crop # with open(os.path.join(right_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: # pickle.dump(right_crop_info, f) cv2.imwrite(os.path.join(right_hand_crop_dir, camera_id + '_' +frame_id + '.png'), right_crop) # 以下为左手 mask_left_hand_idx = np.nonzero(mask_left_hand) if len(mask_left_hand_idx[0]) != 0: left_hand_valid_camera_dict[frame_id].append(camera_id) left_minh = np.min(mask_left_hand_idx[0]) left_maxh = np.max(mask_left_hand_idx[0]) left_minw = np.min(mask_left_hand_idx[1]) left_maxw = np.max(mask_left_hand_idx[1]) left_midh = (left_maxh + left_minh) // 2 left_midw = (left_maxw + left_minw) // 2 left_meanh = np.mean(mask_left_hand_idx[0]) left_meanw = np.mean(mask_left_hand_idx[1]) left_length = max(left_maxh - left_minh, left_maxw - left_minw) half_length = left_length // 2 left_min_y = max(0, left_midh - half_length - margin_size) left_max_y = min(MAX_HEIGHT, left_midh + half_length + margin_size) left_min_x = max(0, left_midw - half_length - margin_size) left_max_x = min(MAX_WIDTH, left_midw + half_length + margin_size) # 将右手变成白色 left_crop = img.copy() left_crop[mask_right_hand == 1] = 255.0 # 居中 left_crop = left_crop[left_min_y: left_max_y, left_min_x: left_max_x] left_crop_info = [left_min_y, left_max_y, left_min_x, left_max_x, left_meanh, left_meanw] if frame_id not in left_crop_info_batch: left_crop_info_batch[frame_id] = {} left_crop_info_batch[frame_id][camera_id] = left_crop_info if frame_id not in left_crop_batch: left_crop_batch[frame_id] = {} left_crop_batch[frame_id][camera_id] = left_crop # with open(os.path.join(left_hand_crop_dir, camera_id + '_' + frame_id + '_crop_info.pkl'), 'wb') as f: # pickle.dump(left_crop_info, f) cv2.imwrite(os.path.join(left_hand_crop_dir, camera_id + '_' + frame_id + '.png'), left_crop) # save valid and invalid info right_hand_invalid_camera_dict = {} left_hand_invalid_camera_dict = {} for frame in frame_list: right_hand_invalid_camera_dict[frame] = [camera for camera in camera_list if camera not in right_hand_valid_camera_dict[frame]] left_hand_invalid_camera_dict[frame] = [camera for camera in camera_list if camera not in left_hand_valid_camera_dict[frame]] left_hand_crop_dir = os.path.join(local_root, date, video_id, 'crop_batch', 'left_hand') right_hand_crop_dir = os.path.join(local_root, date, video_id, 'crop_batch', 'right_hand') os.makedirs(left_hand_crop_dir, exist_ok=True) os.makedirs(right_hand_crop_dir, exist_ok=True) with open(os.path.join(right_hand_crop_dir, represent_frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(right_crop_info_batch, f) with open(os.path.join(right_hand_crop_dir, represent_frame_id + '_crop.pkl'), 'wb') as f: pickle.dump(right_crop_batch, f) with open(os.path.join(left_hand_crop_dir, represent_frame_id + '_crop_info.pkl'), 'wb') as f: pickle.dump(left_crop_info_batch, f) with open(os.path.join(left_hand_crop_dir, represent_frame_id + '_crop.pkl'), 'wb') as f: pickle.dump(left_crop_batch, f) # 寻思这log也没人看啊 # log_path = os.path.join(root, video_id, 'crop', 'log.json') # if os.path.exists(log_path): # with open(log_path, 'r') as f: # log = json.load(f) # log['right_hand_valid_camera_dict'].update(right_hand_valid_camera_dict) # log['left_hand_valid_camera_dict'].update(left_hand_valid_camera_dict) # log['right_hand_invalid_camera_dict'].update(right_hand_invalid_camera_dict) # log['left_hand_invalid_camera_dict'].update(left_hand_invalid_camera_dict) # else: # log = {} # log['right_hand_valid_camera_dict'] = right_hand_valid_camera_dict # log['left_hand_valid_camera_dict'] = left_hand_valid_camera_dict # log['right_hand_invalid_camera_dict'] = right_hand_invalid_camera_dict # log['left_hand_invalid_camera_dict'] = left_hand_invalid_camera_dict # with open(log_path, 'w') as f: # json.dump(log, f) return right_hand_valid_camera_dict, left_hand_valid_camera_dict, right_hand_invalid_camera_dict, left_hand_invalid_camera_dict if __name__ == "__main__": camera_list = ['22070938', '22139905', '22139909', '22139910', '22139911', '22139913', '22139916', '22139946'] # camera_list = ['22139911'] parser = argparse.ArgumentParser() parser.add_argument('--video_id', required=True, type=str) args = parser.parse_args() video_id = args.video_id root = '/share/datasets/HOI-mocap' date = video_id[:8] img_dir = osp.join(root, date, video_id, 'imgs', camera_list[0]) # 默认每个视角的frame数相同 assert os.path.isdir(img_dir) img_filename_list = list(scandir(img_dir, 'png')) frame_list = [] for img_filename in img_filename_list: frame_id = img_filename[-9:-4] frame_list.append(frame_id) procs = [] for camera_id in camera_list: args = (video_id, [camera_id], frame_list) proc = mlp.Process(target=crop_from_mask, args=args) proc.start() procs.append(proc) for i in range(len(procs)): procs[i].join() if __name__ == '__main__': camera_list = ['22070938', '22139905', '22139909', '22139910', '22139911', '22139916', '22139946'] frame_list = [str(i).zfill(5) for i in range(1,3)] crop_from_init('20230904_01', camera_list, frame_list)