File size: 8,549 Bytes
2fd6166
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
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',
}