dataset/behave_paths.py

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

    @staticmethod
    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

    @staticmethod
    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

    @staticmethod
    def get_train_test_from_pkl(pkl_file):
        data = pkl.load(open(pkl_file, 'rb'))
        return data['train'], data['test']

    @staticmethod
    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

    @staticmethod
    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()

    @staticmethod
    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

    @staticmethod
    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

    @staticmethod
    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

    @staticmethod
    def rgb2seq_frame(rgb_file:str):
        "rgb file to seq_name, frame time"
        ss = rgb_file.split(os.sep)
        return ss[-3], ss[-2]

    @staticmethod
    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

    @staticmethod
    def get_seq_name(rgb_file):
        return osp.basename(osp.dirname(osp.dirname(rgb_file)))

    @staticmethod
    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

    @staticmethod
    def rgb2object_name(rgb_file):
        seq_name = DataPaths.get_seq_name(rgb_file)
        obj_name = seq_name.split('_')[2]
        return obj_name

    @staticmethod
    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)

    @staticmethod
    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]

    @staticmethod
    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])

    @staticmethod
    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',
}