Spaces:
Sleeping
Sleeping
| import glob | |
| import os, re | |
| import pickle as pkl | |
| from os.path import join, basename, dirname, isfile | |
| import os.path as osp | |
| import cv2, json | |
| import numpy as np | |
| # PROCESSED_PATH = paths['PROCESSED_PATH'] | |
| BEHAVE_PATH = "/BS/xxie-5/static00/behave_release/sequences/" | |
| RECON_PATH = "/BS/xxie-5/static00/behave-train" | |
| class DataPaths: | |
| """ | |
| class to handle path operations based on BEHAVE dataset structure | |
| """ | |
| def __init__(self): | |
| pass | |
| def load_splits(split_file, dataset_path=None): | |
| assert os.path.exists(dataset_path), f'the given dataset path {dataset_path} does not exist, please check if your training data are placed over there!' | |
| train, val = DataPaths.get_train_test_from_pkl(split_file) | |
| return train, val | |
| # print(train[:5], val[:5]) | |
| if isinstance(train[0], list): | |
| # video data | |
| train_full = [[join(dataset_path, seq[x]) for x in range(len(seq))] for seq in train] | |
| val_full = [[join(dataset_path, seq[x]) for x in range(len(seq))] for seq in val] | |
| else: | |
| train_full = [join(dataset_path, x) for x in train] # full path to the training data | |
| val_full = [join(dataset_path, x) for x in val] # full path to the validation data files | |
| # print(train_full[:5], val_full[:5]) | |
| return train_full, val_full | |
| def load_splits_online(split_file, dataset_path=BEHAVE_PATH): | |
| "load rgb file, smpl and object mesh paths" | |
| keys = ['rgb', 'smpl', 'obj'] | |
| types = ['train', 'val'] | |
| splits = {} | |
| data = pkl.load(open(split_file, 'rb')) | |
| for type in types: | |
| for key in keys: | |
| k = f'{type}_{key}' | |
| splits[k] = [join(dataset_path, x) for x in data[k]] | |
| return splits | |
| def get_train_test_from_pkl(pkl_file): | |
| data = pkl.load(open(pkl_file, 'rb')) | |
| return data['train'], data['test'] | |
| def get_image_paths_seq(seq, tid=1, check_occlusion=False, pat='t*.000'): | |
| """ | |
| find all image paths in one sequence | |
| :param seq: path to one behave sequence | |
| :param tid: test on images from which camera | |
| :param check_occlusion: whether to load full object mask and check occlusion ratio | |
| :return: a list of paths to test image files | |
| """ | |
| image_files = sorted(glob.glob(seq + f"/{pat}/k{tid}.color.jpg")) | |
| # print(image_files, seq + f"/{pat}/k{tid}.color.jpg") | |
| if not check_occlusion: | |
| return image_files | |
| # check object occlusion ratio | |
| valid_files = [] | |
| count = 0 | |
| for img_file in image_files: | |
| mask_file = img_file.replace('.color.jpg', '.obj_rend_mask.png') | |
| if not os.path.isfile(mask_file): | |
| mask_file = img_file.replace('.color.jpg', '.obj_rend_mask.jpg') | |
| full_mask_file = img_file.replace('.color.jpg', '.obj_rend_full.png') | |
| if not os.path.isfile(full_mask_file): | |
| full_mask_file = img_file.replace('.color.jpg', '.obj_rend_full.jpg') | |
| if not isfile(mask_file) or not isfile(full_mask_file): | |
| continue | |
| mask = np.sum(cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE) > 127) | |
| mask_full = np.sum(cv2.imread(full_mask_file, cv2.IMREAD_GRAYSCALE) > 127) | |
| if mask_full == 0: | |
| count += 1 | |
| continue | |
| ratio = mask / mask_full | |
| if ratio > 0.3: | |
| valid_files.append(img_file) | |
| else: | |
| count += 1 | |
| print(f'{mask_file} occluded by {1 - ratio}!') | |
| return valid_files | |
| def get_kinect_id(rgb_file): | |
| "extract kinect id from the rgb file" | |
| filename = osp.basename(rgb_file) | |
| try: | |
| kid = int(filename.split('.')[0][1]) | |
| assert kid in [0, 1, 2, 3, 4, 5], f'found invalid kinect id {kid} for file {rgb_file}' | |
| return kid | |
| except Exception as e: | |
| print(rgb_file) | |
| raise ValueError() | |
| def get_seq_date(rgb_file): | |
| "date for the sequence" | |
| seq_name = str(rgb_file).split(os.sep)[-3] | |
| date = seq_name.split('_')[0] | |
| assert date in ['Date01', 'Date02', 'Date03', 'Date04', 'Date05', 'Date06', 'Date07', | |
| "ICapS01", "ICapS02", "ICapS03", "Date08", "Date09"], f"invalid date for {rgb_file}" | |
| return date | |
| def rgb2obj_path(rgb_file:str, save_name='fit01-smooth'): | |
| "convert an rgb file to a obj mesh file" | |
| ss = rgb_file.split(os.sep) | |
| seq_name = ss[-3] | |
| obj_name = seq_name.split('_')[2] | |
| real_name = obj_name | |
| if 'chair' in obj_name: | |
| real_name = 'chair' | |
| if 'ball' in obj_name: | |
| real_name = 'sports ball' | |
| frame_folder = osp.dirname(rgb_file) | |
| mesh_file = osp.join(frame_folder, real_name, save_name, f'{real_name}_fit.ply') | |
| if not osp.isfile(mesh_file): | |
| # synthetic data | |
| mesh_file = osp.join(frame_folder, obj_name, save_name, f'{obj_name}_fit.ply') | |
| return mesh_file | |
| def rgb2smpl_path(rgb_file:str, save_name='fit03'): | |
| frame_folder = osp.dirname(rgb_file) | |
| real_name = 'person' | |
| mesh_file = osp.join(frame_folder, real_name, save_name, f'{real_name}_fit.ply') | |
| return mesh_file | |
| def rgb2seq_frame(rgb_file:str): | |
| "rgb file to seq_name, frame time" | |
| ss = rgb_file.split(os.sep) | |
| return ss[-3], ss[-2] | |
| def rgb2recon_folder(rgb_file, save_name, recon_path): | |
| "convert rgb file to the subfolder" | |
| dataset_path = osp.dirname(osp.dirname(osp.dirname(rgb_file))) | |
| recon_folder = osp.join(osp.dirname(rgb_file.replace(dataset_path, recon_path)), save_name) | |
| return recon_folder | |
| def get_seq_name(rgb_file): | |
| return osp.basename(osp.dirname(osp.dirname(rgb_file))) | |
| def rgb2template_path(rgb_file): | |
| "return the path to the object template" | |
| from recon.opt_utils import get_template_path | |
| # seq_name = DataPaths.get_seq_name(rgb_file) | |
| # obj_name = seq_name.split('_')[2] | |
| obj_name = DataPaths.rgb2object_name(rgb_file) | |
| path = get_template_path(BEHAVE_PATH+"/../objects", obj_name) | |
| return path | |
| def rgb2object_name(rgb_file): | |
| seq_name = DataPaths.get_seq_name(rgb_file) | |
| obj_name = seq_name.split('_')[2] | |
| return obj_name | |
| def rgb2recon_frame(rgb_file, recon_path=RECON_PATH): | |
| "return the frame folder in recon path" | |
| ss = rgb_file.split(os.sep) | |
| seq_name, frame = ss[-3], ss[-2] | |
| return osp.join(recon_path, seq_name, frame) | |
| def rgb2gender(rgb_file): | |
| "find the gender of this image" | |
| seq_name = str(rgb_file).split(os.sep)[-3] | |
| sub = seq_name.split('_')[1] | |
| return _sub_gender[sub] | |
| def get_dataset_root(rgb_file): | |
| "return the root path to all sequences" | |
| from pathlib import Path | |
| path = Path(rgb_file) | |
| return str(path.parents[2]) | |
| def seqname2gender(seq_name:str): | |
| sub = seq_name.split('_')[1] | |
| return _sub_gender[sub] | |
| ICAP_PATH = "/BS/xxie-6/static00/InterCap" # assume same root folder | |
| date_seqs = { | |
| "Date01": BEHAVE_PATH + "/Date01_Sub01_backpack_back", | |
| "Date02": BEHAVE_PATH + "/Date02_Sub02_backpack_back", | |
| "Date03": BEHAVE_PATH + "/Date03_Sub03_backpack_back", | |
| "Date04": BEHAVE_PATH + "/Date04_Sub05_backpack", | |
| "Date05": BEHAVE_PATH + "/Date05_Sub05_backpack", | |
| "Date06": BEHAVE_PATH + "/Date06_Sub07_backpack_back", | |
| "Date07": BEHAVE_PATH + "/Date07_Sub04_backpack_back", | |
| # "Date08": "/BS/xxie-6/static00/synthesize/Date08_Subxx_chairwood_synzv2-02", | |
| "Date08": "/BS/xxie-6/static00/synz-backup/Date08_Subxx_chairwood_synzv2-02", | |
| "Date09": "/BS/xxie-6/static00/synthesize/Date09_Subxx_obj01_icap", # InterCap sequence synz | |
| "ICapS01": ICAP_PATH + "/ICapS01_sub01_obj01_Seg_0", | |
| "ICapS02": ICAP_PATH + "/ICapS02_sub01_obj08_Seg_0", | |
| "ICapS03": ICAP_PATH + "/ICapS03_sub07_obj05_Seg_0", | |
| } | |
| _sub_gender = { | |
| "Sub01": 'male', | |
| "Sub02": 'male', | |
| "Sub03": 'male', | |
| "Sub04": 'male', | |
| "Sub05": 'male', | |
| "Sub06": 'female', | |
| "Sub07": 'female', | |
| "Sub08": 'female', | |
| } |