src/stacked_hourglass/datasets/sketchfab.py

import os
import glob
import csv
import numpy as np
import cv2
import math
import glob
import pickle as pkl
import open3d as o3d
import trimesh
import torch
import torch.utils.data as data

import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
from configs.anipose_data_info import COMPLETE_DATA_INFO        
from stacked_hourglass.utils.imutils import load_image 
from stacked_hourglass.utils.transforms import crop, color_normalize
from stacked_hourglass.utils.pilutil import imresize 
from stacked_hourglass.utils.imutils import im_to_torch
from configs.dataset_path_configs import TEST_IMAGE_CROP_ROOT_DIR
from configs.data_info import COMPLETE_DATA_INFO_24


class SketchfabScans(data.Dataset):
    DATA_INFO = COMPLETE_DATA_INFO_24
    ACC_JOINTS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 16]  

    def __init__(self, img_crop_folder='default', image_path=None, is_train=False, inp_res=256, out_res=64, sigma=1,
                 scale_factor=0.25, rot_factor=30, label_type='Gaussian', 
                 do_augment='default', shorten_dataset_to=None, dataset_mode='keyp_only'):
        assert is_train == False
        assert do_augment == 'default' or do_augment == False
        self.inp_res = inp_res

        self.n_pcpoints = 3000
        self.folder_imgs = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'datasets', 'sketchfab_test_set', 'images')
        self.folder_silh = self.folder_imgs.replace('images', 'silhouettes')
        self.folder_point_clouds = self.folder_imgs.replace('images', 'point_clouds_' + str(self.n_pcpoints))
        self.folder_meshes = self.folder_imgs.replace('images', 'meshes')
        self.csv_keyp_annots_path = self.folder_imgs.replace('images', 'keypoint_annotations/sketchfab_joint_annotations_complete.csv')
        self.pkl_keyp_annots_path = self.folder_imgs.replace('images', 'keypoint_annotations/sketchfab_joint_annotations_complete_but_as_pkl_file.pkl')
        self.all_mesh_paths = glob.glob(self.folder_meshes + '/**/*.obj', recursive=True)
        name_list = glob.glob(os.path.join(self.folder_imgs, '*.png')) + glob.glob(os.path.join(self.folder_imgs, '*.jpg')) + glob.glob(os.path.join(self.folder_imgs, '*.jpeg'))
        name_list = sorted(name_list)
        # self.test_name_list = [name.split('/')[-1] for name in name_list]
        self.test_name_list = []
        for name in name_list:
            # if not (('13' in name) or ('dalmatian' in name and '1281' in name)):
            # if not ('13' in name):
            self.test_name_list.append(name.split('/')[-1])


        print('len(dataset): ' + str(self.__len__()))
        
        '''
        self.test_mesh_path_list = []
        for img_name in self.test_name_list:
            breed = img_name.split('_')[0]      # will be french instead of french_bulldog
            mask = img_name.split('_')[-2]
            this_mp = []
            for mp in self.all_mesh_paths:
                if (breed in mp) and (mask in mp):
                    this_mp.append(mp)
            if breed in 'french_bulldog':
                this_mp_old = this_mp.copy()
                this_mp = []
                for mp in this_mp_old:
                    if ('_' + mask + '.') in mp:
                        this_mp.append(mp)
            if not len(this_mp) == 1:
                print(breed)
                print(mask)
                this_mp[0].index(mask)
                import pdb; pdb.set_trace()
            else:
                self.test_mesh_path_list.append(this_mp[0])

        all_pc_paths = []
        for index in range(len(self.test_name_list)):
            img_name = self.test_name_list[index]
            dog_name = img_name.split('_' + img_name.split('_')[-1])[0]
            breed = img_name.split('_')[0]      # will be french instead of french_bulldog
            mask = img_name.split('_')[-2]
            path_pc = self.folder_point_clouds + '/' + dog_name + '.ply'
            if not path_pc in all_pc_paths:
                try:
                    print(path_pc)
                    mesh_path = self.test_mesh_path_list[index]
                    mesh_gt = o3d.io.read_triangle_mesh(mesh_path)
                    n_points = 3000     # 20000
                    pointcloud = mesh_gt.sample_points_uniformly(number_of_points=n_points)
                    o3d.io.write_point_cloud(path_pc, pointcloud, write_ascii=False, compressed=False, print_progress=False)
                    all_pc_paths.append(path_pc)
                except:
                    print(path_pc)
        '''

        # import pdb; pdb.set_trace()

        self.test_mesh_path_list = []
        self.all_pc_paths = []
        for index in range(len(self.test_name_list)):
            img_name = self.test_name_list[index]
            dog_name = img_name.split('_' + img_name.split('_')[-1])[0]
            breed = img_name.split('_')[0]      # will be french instead of french_bulldog
            mask = img_name.split('_')[-2]
            mesh_path = self.folder_meshes + '/' + dog_name + '.obj'
            path_pc = self.folder_point_clouds + '/' + dog_name + '.ply'
            if dog_name in  ['dalmatian_1281', 'french_bulldog_13']:
                # mesh_path_for_pc = '/is/cluster/work/nrueegg/icon_pifu_related/barc_for_bite/datasets/sketchfab_test_set/meshes_old/dalmatian/1281/Renderbot-animal-obj-1281.obj'
                mesh_path_for_pc = self.folder_meshes + '/' + dog_name + '_simple.obj'
            else:
                mesh_path_for_pc = mesh_path 
            self.test_mesh_path_list.append(mesh_path)
            # if not path_pc in self.all_pc_paths:
            if os.path.isfile(path_pc):
                self.all_pc_paths.append(path_pc)
            else:
                try:
                    mesh_gt = o3d.io.read_triangle_mesh(mesh_path_for_pc)
                except:
                    import pdb; pdb.set_trace()
                    mesh = trimesh.load(mesh_path_for_pc, process=False,  maintain_order=True) 
                    vertices = mesh.vertices
                    faces = mesh.faces

                print(mesh_path_for_pc)
                pointcloud = mesh_gt.sample_points_uniformly(number_of_points=self.n_pcpoints)
                o3d.io.write_point_cloud(path_pc, pointcloud, write_ascii=False, compressed=False, print_progress=False)
                self.all_pc_paths.append(path_pc)
                # except:
                #     print(path_pc)

        # add keypoint annotations (mesh vertices)
        read_annots_from_csv = False        # True
        if read_annots_from_csv:
            self.all_keypoint_annotations, self.keypoint_name_dict = self._read_keypoint_csv(self.csv_keyp_annots_path, folder_meshes=self.folder_meshes, get_keyp_coords=True)
            with open(self.pkl_keyp_annots_path, 'wb') as handle:
                pkl.dump(self.all_keypoint_annotations, handle, protocol=pkl.HIGHEST_PROTOCOL)
        else:
            with open(self.pkl_keyp_annots_path, 'rb') as handle:
                self.all_keypoint_annotations = pkl.load(handle)





    def _read_keypoint_csv(self, csv_path, folder_meshes=None, get_keyp_coords=True, visualize=False):
        with open(csv_path,'r') as f:
            reader = csv.reader(f)
            headers = next(reader)
            row_list = [{h:x for (h,x) in zip(headers,row)} for row in reader]
            assert(headers[2] == 'hiwi')
        keypoint_names = headers[3:]
        center_keypoint_names = ['nose','tail_start','tail_end']
        right_keypoint_names = ['right_front_paw','right_front_elbow','right_back_paw','right_back_hock','right_ear_top','right_ear_bottom','right_eye']
        left_keypoint_names = ['left_front_paw','left_front_elbow','left_back_paw','left_back_hock','left_ear_top','left_ear_bottom','left_eye']
        keypoint_name_dict = {'all': keypoint_names, 'left': left_keypoint_names, 'right': right_keypoint_names, 'center': center_keypoint_names}
        # prepare output dicts
        all_keypoint_annotations = {}
        for ind in range(len(row_list)):
            name = row_list[ind]['mesh_name']
            this_dict = row_list[ind]
            del this_dict['hiwi']
            all_keypoint_annotations[name] = this_dict
            keypoint_idxs = np.zeros((len(keypoint_names), 2))
            if get_keyp_coords:
                mesh_path = folder_meshes + '/' + row_list[ind]['mesh_name']
                mesh = trimesh.load(mesh_path, process=False,  maintain_order=True) 
                vertices = mesh.vertices
                keypoint_3d_locations = np.zeros((len(keypoint_names), 4))      # 1, 2, 3: coords, 4: is_valid
            for ind_kp, name_kp in enumerate(keypoint_names):
                idx = this_dict[name_kp]
                if idx in ['', 'n/a']:
                    keypoint_idxs[ind_kp, 0] = -1
                else:
                    keypoint_idxs[ind_kp, 0] = this_dict[name_kp]
                    keypoint_idxs[ind_kp, 1] = 1        # is valid
                    if get_keyp_coords:
                        keyp = vertices[int(row_list[ind][name_kp])]
                        keypoint_3d_locations[ind_kp, :3] = keyp
                        keypoint_3d_locations[ind_kp, 3] = 1
            all_keypoint_annotations[name]['all_keypoint_vertex_idxs'] = keypoint_idxs
            if get_keyp_coords:
                all_keypoint_annotations[name]['all_keypoint_coords_and_isvalid'] = keypoint_3d_locations
        # create visualizations if desired
        if visualize:
            raise NotImplementedError       # only debug path is missing
            out_path = '.... some debug path'
            red_color = np.asarray([255, 0, 0], dtype=np.uint8)
            green_color = np.asarray([0, 255, 0], dtype=np.uint8)
            blue_color = np.asarray([0, 0, 255], dtype=np.uint8)
            for ind in range(len(row_list)):
                mesh_path = folder_meshes + '/' + row_list[ind]['mesh_name']
                mesh = trimesh.load(mesh_path, process=False,  maintain_order=True)         # maintain_order is very important!!!!!
                vertices = mesh.vertices
                faces = mesh.faces
                dog_mesh_nocolor = trimesh.Trimesh(vertices=vertices, faces=faces, process=False, maintain_order=True)
                dog_mesh_nocolor.visual.vertex_colors = np.ones_like(vertices, dtype=np.uint8) * 255
                sphere_list = [dog_mesh_nocolor]
                for keyp_name in keypoint_names:
                    if not (row_list[ind][keyp_name] == '' or row_list[ind][keyp_name] == 'n/a'):
                        keyp = vertices[int(row_list[ind][keyp_name])]
                        sphere = trimesh.primitives.Sphere(radius=0.02, center=keyp)
                        if keyp_name in right_keypoint_names:
                            colors = np.ones_like(sphere.vertices) * red_color[None, :]
                        elif keyp_name in left_keypoint_names:
                            colors = np.ones_like(sphere.vertices) * blue_color[None, :]
                        else:
                            colors = np.ones_like(sphere.vertices) * green_color[None, :]
                        sphere.visual.vertex_colors = colors  # trimesh.visual.random_color()
                        sphere_list.append(sphere)
                scene_keyp = trimesh.Scene(sphere_list)
                scene_keyp.export(out_path + os.path.basename(mesh_path).replace('.obj', '_withkeyp.obj'))
        return all_keypoint_annotations, keypoint_name_dict



    def __getitem__(self, index):
        img_name = self.test_name_list[index]
        dog_name = img_name.split('_' + img_name.split('_')[-1])[0]
        breed = img_name.split('_')[0]      # will be french instead of french_bulldog
        mask = img_name.split('_')[-2]
        mesh_path = self.test_mesh_path_list[index]
        # mesh_gt = o3d.io.read_triangle_mesh(mesh_path)

        path_pc = self.folder_point_clouds + '/' + dog_name + '.ply'
        assert path_pc in self.all_pc_paths
        pc_trimesh = trimesh.load(path_pc, process=False, maintain_order=True)
        pc_points = np.asarray(pc_trimesh.vertices)
        assert pc_points.shape[0] == self.n_pcpoints


        # get annotated 3d keypoints
        keyp_3d = self.all_keypoint_annotations[mesh_path.split('/')[-1]]['all_keypoint_coords_and_isvalid']


        # load image
        img_path = os.path.join(self.folder_imgs, img_name)
        
        img = load_image(img_path)  # CxHxW
        # try on silhouette images!
        # seg_path = os.path.join(self.folder_silh, img_name)
        # img = load_image(seg_path)  # CxHxW

        img_vis = np.transpose(img, (1, 2, 0))
        seg_path = os.path.join(self.folder_silh, img_name)
        seg = cv2.imread(seg_path, cv2.IMREAD_UNCHANGED)[:, :, 3]
        seg[seg>0] = 1
        seg_s0 = np.nonzero(seg.sum(axis=1)>0)[0] 
        seg_s1 = np.nonzero(seg.sum(axis=0)>0)[0] 
        bbox_xywh = [seg_s1.min(), seg_s0.min(), seg_s1.max() - seg_s1.min(), seg_s0.max() - seg_s0.min()]
        bbox_c = [bbox_xywh[0]+0.5*bbox_xywh[2], bbox_xywh[1]+0.5*bbox_xywh[3]]
        bbox_max = max(bbox_xywh[2], bbox_xywh[3])
        bbox_diag = math.sqrt(bbox_xywh[2]**2 + bbox_xywh[3]**2)
        # bbox_s = bbox_max / 200.      # the dog will fill the image -> bbox_max = 256
        # bbox_s = bbox_diag / 200.     # diagonal of the boundingbox will be 200
        bbox_s = bbox_max / 200. * 256. / 200.  # maximum side of the bbox will be 200
        c = torch.Tensor(bbox_c)
        s = bbox_s
        r = 0

        # Prepare image and groundtruth map
        inp_col = crop(img, c, s, [self.inp_res, self.inp_res], rot=r)
        inp = color_normalize(inp_col, self.DATA_INFO.rgb_mean, self.DATA_INFO.rgb_stddev)

        silh_3channels = np.stack((seg, seg, seg), axis=0)
        inp_silh = crop(silh_3channels, c, s, [self.inp_res, self.inp_res], rot=r)

        '''
        # prepare image (cropping and color)
        img_max = max(img.shape[1], img.shape[2])
        img_padded = torch.zeros((img.shape[0], img_max, img_max))
        if img_max == img.shape[2]:
            start = (img_max-img.shape[1])//2
            img_padded[:, start:start+img.shape[1], :] = img
        else:
            start = (img_max-img.shape[2])//2
            img_padded[:, :, start:start+img.shape[2]] = img   
        img = img_padded
        img_prep = im_to_torch(imresize(img, [self.inp_res, self.inp_res], interp='bilinear'))   
        inp = color_normalize(img_prep, self.DATA_INFO.rgb_mean, self.DATA_INFO.rgb_stddev)
        '''
        # add the following fields to make it compatible with stanext, most of them are fake
        target_dict = {'index': index, 'center' : -2, 'scale' : -2, 
            'breed_index': -2, 'sim_breed_index': -2,
            'ind_dataset': 1}
        target_dict['pts'] = np.zeros((self.DATA_INFO.n_keyp, 3))
        target_dict['tpts'] = np.zeros((self.DATA_INFO.n_keyp, 3))
        target_dict['target_weight'] = np.zeros((self.DATA_INFO.n_keyp, 1))
        target_dict['silh'] = inp_silh[0, :, :]      # np.zeros((self.inp_res, self.inp_res))
        target_dict['mesh_path'] = mesh_path
        target_dict['pointcloud_path'] = path_pc
        target_dict['pointcloud_points'] = pc_points
        target_dict['keypoints_3d'] = keyp_3d
        return inp, target_dict


    def __len__(self):
        return len(self.test_name_list)