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)