def _gen_sign(cls, params, ret):
def join(format, params, offset, sep='\n'):
return sep.join(cls.SIGN_PADDING + format(*p) for p in params)[offset:]
ret_stype, _, ret_desc = ret
ret = [(ret_stype, ret_desc)]
ret_offset = len(cls.SIGN_PADDING) - len('Parameters') + len('Return')
return align(cls.SIGNATURE.format(
params=join(cls.PARAM_TEMPLATE.format, params, len(cls.SIGN_PADDING)),
ret=join('{0} - {1}'.format, ret, ret_offset)))
def VirtualFree(self, uc, esp, log, address, size, free_type):
log and print(
f"VirtualFree: chunk to free: 0x{address:02x}, size 0x{size:02x}, type 0x{free_type:02x}"
)
new_chunks = []
success = False
for start, end in sorted(self.allocated_chunks):
if start <= address <= end:
if free_type & 0x8000 and size == 0: # MEM_RELEASE, clear whole allocated range
if address in self.alloc_sizes:
end_addr = self.alloc_sizes[address]
uc.mem_unmap(address, align(size))
new_chunks += remove_range((start, end),
(address, end_addr))
success = True
else:
log and print(
f"\t0x{address} is not an alloc base address!")
new_chunks += [(start, end)]
elif free_type & 0x4000 and size > 0: # MEM_DECOMMIT, free requested size
end_addr = address + align(size)
uc.mem_unmap(address, align(size))
new_chunks += remove_range((start, end),
(address, end_addr))
success = True
else:
log and print("\tIncorrect size + type combination!")
new_chunks += [(start, end)]
else:
new_chunks += [(start, end)]
self.allocated_chunks = list(merge(new_chunks))
log and self.print_allocs()
if success:
return 1
log and print("\tAddress range not allocated!")
return 0
def duplicate(img, copies_to_img_x, copies_to_img_y):
current = img.copy()
current_to_img_x = 0
current_to_img_y = 0
for copy_to_img_x, copy_to_img_y in zip(copies_to_img_x, copies_to_img_y):
copy_to_current_x = copy_to_img_x - current_to_img_x
copy_to_current_y = copy_to_img_y - current_to_img_y
copy_aligned, current_aligned, aligned_to_current_x, aligned_to_current_y = utils.align(
img, current, copy_to_current_x, copy_to_current_y)
current = np.logical_or(copy_aligned, current_aligned)
current_to_img_x = aligned_to_current_x + current_to_img_x
current_to_img_y = aligned_to_current_y + current_to_img_y
return utils.trim_binary_image(current)
def get_hsps(seq1_dict, seq2_dict, K, scoring_matrix, T):
hsps = []
for key, val in seq1_dict.items():
neighbors = utils.find_neighbors(key, scoring_matrix, ALPHABET, T)
for neighbor in neighbors:
if neighbor in seq2_dict.keys():
score = utils.align(key, neighbor, scoring_matrix)
for seq1_start in val:
seq1_end = seq1_start + K - 1
for seq2_start in seq2_dict[neighbor]:
seq2_end = seq2_start + K - 1
hsp = HSP(seq1_start, seq1_end, seq2_start, seq2_end,
score)
hsps.append(hsp)
return hsps
def main():
args = get_args()
img1, img2 = align(*args.image)
bg1, bg2 = args.background
bg_template = np.ones_like(img1).astype("float")
if args.reverse:
img1, img2, bg1, bg2 = img2, img1, bg2, bg1
res, _ = get_mixed(
norm_img(img1),
norm_img(img2),
norm_img(bg1) * bg_template,
norm_img(bg2) * bg_template,
args.adjust_ratio,
)
cv.imwrite(args.output, res)
def alloc(self, log, size, uc, offset=None):
page_size = 4 * 1024
aligned_size = align(size, page_size)
log and print(
f"\tUnaligned size: 0x{size:02x}, aligned size: 0x{aligned_size:02x}"
)
if offset is None:
for chunk_start, chunk_end in self.allocated_chunks:
if chunk_start <= self.dynamic_mem_offset <= chunk_end:
# we have to push back the dynamic mem offset as it is inside an already allocated chunk!
self.dynamic_mem_offset = chunk_end + 1
offset = self.dynamic_mem_offset
self.dynamic_mem_offset += aligned_size
new_offset_m = offset % page_size
aligned_address = offset - new_offset_m
# check if we have mapped parts of it already
mapped_partial = False
for chunk_start, chunk_end in self.allocated_chunks:
if chunk_start <= aligned_address < chunk_end:
if aligned_address + aligned_size <= chunk_end:
log and print(f"\tAlready fully mapped")
else:
log and print(
f"\tMapping missing piece 0x{chunk_end + 1:02x} to 0x{aligned_address + aligned_size:02x}"
)
uc.mem_map(chunk_end,
aligned_address + aligned_size - chunk_end)
mapped_partial = True
break
if not mapped_partial:
uc.mem_map(aligned_address, aligned_size)
log and print(
f"\tfrom 0x{aligned_address:02x} to 0x{(aligned_address + aligned_size):02x}"
)
self.allocated_chunks = list(
merge(self.allocated_chunks +
[(aligned_address, aligned_address + aligned_size)]))
log and self.print_allocs()
self.alloc_sizes[aligned_address] = aligned_size
return aligned_address
def xor_diff(img, diff_to_ref_x, diff_to_ref_y, diff, ref):
"""
:param img:
:param diff_to_ref_x:
:param diff_to_ref_y:
:param diff:
:param ref:
:return:
"""
if 0 == diff.sum():
return img
_, img_to_ref_x, img_to_ref_y = jaccard.jaccard_coef(img, ref)
diff_to_img_x = diff_to_ref_x - img_to_ref_x
diff_to_img_y = diff_to_ref_y - img_to_ref_y
diff_aligned, img_aligned, _, _ = utils.align(diff, img, diff_to_img_x,
diff_to_img_y)
return utils.erase_noise_point(
utils.trim_binary_image(np.logical_xor(img_aligned, diff_aligned)), 4)
def apply_binary_transformation(imgA,
imgB,
tran,
imgA_to_imgB_x=None,
imgA_to_imgB_y=None,
imgC=None):
if imgA_to_imgB_x is None or imgA_to_imgB_y is None:
align_foo_name = tran.get("align")
if align_foo_name is None:
_, imgA_to_imgB_x, imgA_to_imgB_y = jaccard.jaccard_coef(
imgA, imgB)
else:
align_foo = getattr(THIS, align_foo_name)
imgA_to_imgB_x, imgA_to_imgB_y = align_foo(imgA, imgB, imgC)
imgA_aligned, imgB_aligned, aligned_to_B_x, aligned_to_B_y = utils.align(
imgA, imgB, imgA_to_imgB_x, imgA_to_imgB_y)
img_aligned = None
binary_trans = tran.get("value")
for tran in binary_trans:
name = tran.get("name")
if name is None:
continue
foo = getattr(THIS, name)
args = tran.get("args")
if args is None:
img_aligned = foo(imgA_aligned, imgB_aligned)
else:
img_aligned = foo(imgA_aligned, imgB_aligned, **args)
return img_aligned, int(imgA_to_imgB_x), int(
imgA_to_imgB_y), aligned_to_B_x, aligned_to_B_y
def __init__(self, exploit, start, size, align_to=None, alignment=None):
self.exploit = exploit
if align_to is not None:
if start < align_to:
raise Exception("Trying to align to a something which is after our buffer: aligning " + self.exploit.pointer_format % start + " to " + self.exploit.pointer_format % align_to)
self.align_to = align_to
self.alignment = size if (alignment is None) else alignment
self.start = align(start, self.align_to, self.alignment)
self.index = (self.start - self.align_to) / self.alignment
else:
self.alignment = 1
self.align_to = 0
self.start = start
self.index = 0
self.content = ""
self.pointer = self.exploit.ptr2str(self.start)
self.size = size
self.end = self.start + self.size
self.wasted = -1
self.is_buffered = False
if self.index < 0:
log("Warning: a negative index has been computed: " + str(self.index))
def subtract_diff(img, diff_to_ref_x, diff_to_ref_y, diff, ref, coords=False):
"""
:param img:
:param diff_to_ref_x:
:param diff_to_ref_y:
:param diff:
:param ref:
:return:
"""
if 0 == diff.sum():
if coords:
return img, 0, 0
else:
return img
_, img_to_ref_x, img_to_ref_y = jaccard.jaccard_coef(img, ref)
diff_to_img_x = diff_to_ref_x - img_to_ref_x
diff_to_img_y = diff_to_ref_y - img_to_ref_y
diff_aligned, img_aligned, aligned_to_img_x, aligned_to_img_y = utils.align(
diff, img, diff_to_img_x, diff_to_img_y)
result, result_to_aligned_x, result_to_aligned_y = utils.trim_binary_image(
utils.erase_noise_point(
np.logical_and(img_aligned, np.logical_not(diff_aligned)), 8),
coord=True)
if coords:
diff_to_result_x = (diff_to_img_x -
aligned_to_img_x) - result_to_aligned_x
diff_to_result_y = (diff_to_img_y -
aligned_to_img_y) - result_to_aligned_y
return result, diff_to_result_x, diff_to_result_y
else:
return result
def start_estimation(self):
n = 3
log_file = os.path.join(self.save_dir, 'error_log.txt')
self.f_log = open(log_file, 'w')
image_id = 0
last_estimation = None
while n > 0:
n = n - 1
dataset_my = MYDataset(self.synset_names, 'val', self.config)
data_path = self.pack_path + '/' + subfile
if not os.path.exists(data_path):
os.makedirs(data_path)
load_data_res = dataset_my.load_data(data_path)
if not load_data_res:
rospy.logerr(
"Load data Failed. Check path of the input images.")
return False
dataset_my.prepare(self.class_map)
dataset = dataset_my
print('*' * 50)
image_start = time.time()
image_path = dataset.image_info[image_id]["path"]
result = {}
image = dataset.load_image(image_id)
depth = dataset.load_depth(image_id)
result['image_id'] = image_id
result['image_path'] = image_path
image_path_parsing = image_path.split('/')
## detection
start = time.time()
with self.graph.as_default():
set_session(self.sess)
detect_result = self.model.detect([image], verbose=0)
r = detect_result[0]
elapsed = time.time() - start
print('\nDetection takes {:03f}s.'.format(elapsed))
result['pred_class_ids'] = r['class_ids']
result['pred_bboxes'] = r['rois']
result['pred_RTs'] = None
result['pred_scores'] = r['scores']
if len(r['class_ids']) == 0:
rospy.logwarn('Nothing detected.')
dataset_my.detect_finished()
continue
print('Aligning predictions...')
start = time.time()
result['pred_RTs'], result[
'pred_scales'], error_message, elapses = utils.align(
r['class_ids'], r['masks'], r['coords'], depth,
self.intrinsics, self.synset_names, image_path)
print('New alignment takes {:03f}s.'.format(time.time() - start))
print("The predicted transfrom matrix:")
print(result['pred_RTs'])
filtered_RTs = RT_filter(result['pred_RTs'], self.detect_range)
if len(filtered_RTs) == 0:
rospy.logwarn('No detected object in %d cm.',
self.detect_range)
dataset_my.detect_finished()
continue
if last_estimation is not None:
pose_diff = position_diff(last_estimation, filtered_RTs[0])
if pose_diff < 0.05:
n = 0
self.concur.set_T(filtered_RTs[0],
self.camera_optical_frame)
if self.first_call:
self.concur.start()
self.first_call = False
else:
self.concur.resume()
else:
rospy.loginfo(
'The pose estimation is not stable. Distance to last estimation: %f m.',
pose_diff)
last_estimation = filtered_RTs[0]
if len(error_message):
self.f_log.write(error_message)
if self.draw:
draw_rgb = False
utils.draw_detections_wogt(
image, self.save_dir, data,
image_path_parsing[-2] + '_' + image_path_parsing[-1],
self.intrinsics, self.synset_names, draw_rgb, r['rois'],
r['class_ids'], r['masks'], r['coords'],
result['pred_RTs'], r['scores'], result['pred_scales'])
path_parse = image_path.split('/')
image_short_path = '_'.join(path_parse[-3:])
save_path = os.path.join(self.save_dir,
'results_{}.pkl'.format(image_short_path))
with open(save_path, 'wb') as f:
cPickle.dump(result, f)
print('Results of image {} has been saved to {}.'.format(
image_short_path, save_path))
elapsed = time.time() - image_start
print('Takes {} to finish this image.'.format(elapsed))
print('Alignment time: ', elapses)
print('\n')
dataset_my.detect_finished()
self.f_log.close()
return True
def save_img_head(frame, save_path, seq, cam, cam_id, json_file, frame_id,
threshold, yaw_ref):
img_path = os.path.join(save_path, seq)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = Image.fromarray(frame)
# print(frame.size)
E_ref = np.mat([[1, 0, 0, 0.], [0, -1, 0, 0], [0, 0, -1, 50], [0, 0, 0,
1]])
cam['K'] = np.mat(cam['K'])
cam['distCoef'] = np.array(cam['distCoef'])
cam['R'] = np.mat(cam['R'])
cam['t'] = np.array(cam['t']).reshape((3, 1))
with open(json_file) as dfile:
fframe = json.load(dfile)
count_face = -1
yaw_avg = 0
for face in fframe['people']:
# 3D Face has 70 3D joints, stored as an array [x1,y1,z1,x2,y2,z2,...]
face3d = np.array(face['face70']['landmarks']).reshape(
(-1, 3)).transpose()
face_conf = np.asarray(face['face70']['averageScore'])
model_points_3D = np.ones((4, 58), dtype=np.float32)
model_points_3D[0:3] = model_points
clean_match = (face_conf[kp_idx] > 0.1
) #only pick points confidence higher than 0.1
kp_idx_clean = kp_idx[clean_match]
kp_idx_model_clean = kp_idx_model[clean_match]
if (len(kp_idx_clean) > 6):
count_face += 1
rotation, translation, error, scale = align(
np.mat(model_points_3D[0:3, kp_idx_model_clean]),
np.mat(face3d[:, kp_idx_clean]))
sphere_new = scale * rotation @ (sphere) + translation
pt_helmet = projectPoints(sphere_new, cam['K'], cam['R'], cam['t'],
cam['distCoef'])
temp = np.zeros((4, 4))
temp[0:3, 0:3] = rotation
temp[0:3, 3:4] = translation
temp[3, 3] = 1
E_virt = np.linalg.inv(temp @ np.linalg.inv(E_ref))
E_real = np.zeros((4, 4))
E_real[0:3, 0:3] = cam['R']
E_real[0:3, 3:4] = cam['t']
E_real[3, 3] = 1
compound = E_real @ np.linalg.inv(E_virt)
status, [pitch, yaw, roll
] = select_euler(np.rad2deg(inverse_rotate_zyx(compound)))
yaw = -yaw
roll = -roll
yaw_avg = yaw_avg + yaw
if (abs(yaw - yaw_ref) > threshold or yaw_ref == -999):
if (status == True):
x_min = int(max(min(pt_helmet[0, :]), 0))
y_min = int(max(min(pt_helmet[1, :]), 0))
x_max = int(min(max(pt_helmet[0, :]), frame.size[0]))
y_max = int(min(max(pt_helmet[1, :]), frame.size[1]))
# print(x_min, y_min, x_max, y_max)
if (x_min < x_max and y_min < y_max and abs(x_min - x_max)
< frame.size[0]): #some sanity check
h = y_max - y_min
w = x_max - x_min
if not (h / w > 2 or w / h > 2
): #eleminate those too wide or too narrow
img = frame.crop((x_min, y_min, x_max, y_max))
# draw = ImageDraw.Draw(img)
# draw.text((0, 10), "yaw: {}".format(round(yaw)), (0, 255, 255))
# draw.text((0, 0), "pitch: {}".format(round(pitch)), (0, 255, 255))
# draw.text((0, 20), "roll: {}".format(round(roll)), (0, 255, 255))
# plt.imshow(img)
# plt.show()
filename = '{0:02d}_{1:01d}_{2:08d}.jpg'.format(
cam_id, count_face, frame_id)
if not (os.path.exists(img_path)):
os.mkdir(img_path)
file_path = os.path.join(img_path, filename)
img.save(file_path, "JPEG")
anno_path = os.path.join(save_path,
"annotation.txt")
line = seq + '/' + filename + ',' + str(
yaw) + ',' + str(pitch) + ',' + str(
roll) + '\n'
with open(anno_path, "a") as f:
f.write(line)
if count_face != -1:
return yaw_avg / (count_face + 1)
else:
return -999
# read jump table header
jt_header = struct.unpack('>LLLL', jt_data[0:16])
# read first entry of the jump table
jt_1st_entry = struct.unpack('>HHHH', jt_data[16:24])
if jt_1st_entry[1] != 0x3F3C or jt_1st_entry[3] != 0xA9F0:
print("Invalid jump table! 1st entry is corrupted!")
exit(1)
mt = MacTraps(rt, args.path)
jt_length = jt_header[2]
jt_offset = jt_header[3]
print("JT length:", hex(jt_offset + jt_length))
jt_handle = mt._mm.new_handle(
utils.align(jt_offset + jt_length, 0x10000))
a5_base = mem.r32(jt_handle)
# copy JT from CODE,0 to A5 + jt_offset
for i in range(jt_length):
mem.w8(a5_base + jt_offset + i, jt_data[16 + i])
ep_seg_id = jt_1st_entry[2]
ep_offset = jt_1st_entry[0]
print("1st entry of the JT points to segment %d, offset %x" % (ep_seg_id, ep_offset))
print("Loading code segment %d" % ep_seg_id)
ep_res = rf[b'CODE'][ep_seg_id]
ep_res_len = ep_res.length
ep_data = ep_res.data_raw
prog_handle = mt._mm.new_handle(ep_res_len)
def manual_data_handling():
# Loading all raw information
print("Loading raw data")
instance_dict = {}
# Loading instance information from meta file
with open(str(meta_path), 'r') as f:
for line in f:
line_info = line.split(' ')
instance_id = int(line_info[0])
class_id = int(line_info[1])
instance_dict[instance_id] = class_id
#print(f"Instance #{instance_id}: {class_id} = {pj.constants.NOCS_CLASSES[class_id]}")
# Loading data
color_image = cv2.imread(str(image_path), cv2.IMREAD_UNCHANGED)
mask_image = cv2.imread(str(mask_path), cv2.IMREAD_UNCHANGED)[:, :, 2]
coord_map = cv2.imread(str(coord_path), cv2.IMREAD_UNCHANGED)[:, :, :3]
coord_map = coord_map[:, :, (2, 1, 0)]
depth_image = cv2.imread(str(depth_path), cv2.IMREAD_UNCHANGED)
# Converting depth to the correct shape and dtype
if len(depth_image.shape) == 3: # encoded depth image
new_depth = np.uint16(depth_image[:, :, 1] * 256) + np.uint16(
depth_image[:, :, 2])
new_depth = new_depth.astype(np.uint16)
depth_image = new_depth
elif len(depth_image.shape) == 2 and depth_image.dtype == 'uint16':
pass # depth is perfecto!
else:
assert False, '[ Error ]: Unsupported depth type'
"""
cv2.imshow('mask_image', mask_image)
cv2.imshow('coord_map', coord_map)
cv2.imshow('depth_image', depth_image)
cv2.waitKey(0)
"""
# Processing data
print("Processing data")
cdata = np.array(mask_image, dtype=np.int32)
# The value of the mask, from 0 to 255, is the class id
instance_ids = list(np.unique(cdata))
instance_ids = sorted(instance_ids)
#print(f'instance_ids: {instance_ids}')
# removing background
assert instance_ids[-1] == 255
del instance_ids[-1]
cdata[cdata == 255] = -1
assert (np.unique(cdata).shape[0] < 20)
num_instance = len(instance_ids)
h, w = cdata.shape
# flip z axis of coord map
coord_map = np.array(coord_map, dtype=np.float32) / 255
coord_map[:, :, 2] = 1 - coord_map[:, :, 2]
masks = np.zeros([h, w, num_instance], dtype=np.uint8)
coords = np.zeros((h, w, num_instance, 3), dtype=np.float32)
class_ids = np.zeros([num_instance], dtype=np.int_)
scales = np.zeros([num_instance, 3], dtype=np.float32)
# Determing the scales (of the 3d bounding boxes)
with open(str(meta_path), 'r') as f:
lines = f.readlines()
scale_factor = np.zeros((len(lines), 3), dtype=np.float32)
for i, line in enumerate(lines):
words = line[:-1].split(' ')
symmetry_id = words[2]
reference_id = words[3]
bbox_file = obj_model_dir / symmetry_id / reference_id / 'bbox.txt'
bbox = np.loadtxt(str(bbox_file))
value = bbox[0, :] - bbox[1, :]
scale_factor[i, :] = value # [a b c] for scale of NOCS [x y z]
# Deleting background objects and non-existing objects
instance_dict = {
instance_id: class_id
for instance_id, class_id in instance_dict.items()
if (class_id != 0 and instance_id in instance_ids)
}
i = 0
for instance_id in instance_ids:
if instance_id not in instance_dict.keys():
continue
instance_mask = np.equal(cdata, instance_id)
assert np.sum(instance_mask) > 0
assert instance_dict[instance_id]
masks[:, :, i] = instance_mask
coords[:, :,
i, :] = np.multiply(coord_map,
np.expand_dims(instance_mask, axis=-1))
# class ids are also one-indexed
class_ids[i] = instance_dict[instance_id]
scales[i, :] = scale_factor[instance_id - 1, :]
i += 1
masks = masks[:, :, :i]
coords = coords[:, :, :i, :]
coords = np.clip(coords, 0, 1) # normalize
class_ids = class_ids[:i]
scales = scales[:i]
bboxes = nocs_utils.extract_bboxes(masks)
scores = [100 for j in range(len(class_ids))]
"""
print(f"masks.shape: {masks.shape}")
print(f"coords.shape: {coords.shape}")
print(f"scales: {scales}")
for i in range(masks.shape[2]):
visual_mask = masks[:,:,i]
visual_mask[visual_mask==0] = 255
cv2.imshow(f"({pj.constants.NOCS_CLASSES[class_ids[i]]}) - Mask {i}", visual_mask)
cv2.imshow(f"({pj.constants.NOCS_CLASSES[class_ids[i]]}) - Coords {i}", coords[:,:,i,:])
cv2.waitKey(0)
"""
# now obtaining the rotation and translation
RTs, _, _, _ = nocs_utils.align(class_ids, masks, coords, depth_image,
pj.constants.INTRINSICS,
pj.constants.NOCS_CLASSES, ".", None)
print(f'color_image.shape: {color_image.shape}')
print(
f'depth_image.shape: {depth_image.shape} depth_image.dtype: {depth_image.dtype}'
)
print(
f'mask_image.shape: {mask_image.shape} mask_image.dtype: {mask_image.dtype}'
)
print(
f'coord_map.shape: {coord_map.shape} coord_map.dtype: {coord_map.dtype}'
)
print(f'class_ids: {class_ids}')
print(f'domain_label: {None}')
print(f'bboxes: {bboxes}')
print(f'scales: {scales}')
print(f'RTs: {RTs}')
return class_ids, bboxes, masks, coords, RTs, scores, scales, instance_dict
if 'handle_visibility' in gt:
result['gt_handle_visibility'] = gt[
'handle_visibility']
assert len(
gt['handle_visibility']) == len(gt_class_ids)
print('got handle visibiity.')
else:
result['gt_handle_visibility'] = np.ones_like(
gt_class_ids)
else: # align gt coord with depth to get RT
if not data in ['coco_val', 'coco_train']:
if len(gt_class_ids) == 0:
print('No gt instance exsits in this image.')
print('\nAligning ground truth...')
result['gt_RTs'], _, error_message, _ = utils.align(
gt_class_ids, gt_mask, gt_coord, depth, intrinsics,
synset_names, image_path,
save_dir + '/' + f'{xx}_gt_')
if len(error_message): f_log.write(error_message)
result['gt_handle_visibility'] = np.ones_like(gt_class_ids)
## detection
start = time.time()
detect_result = model.detect([image], verbose=0)
dr = detect_result[0]
print('\nDetection takes {:03f}s.'.format(time.time() - start))
result['pred_class_ids'] = dr['class_ids']
result['pred_bboxes'] = dr['rois']
result['pred_RTs'] = None
result['pred_scores'] = dr['scores']
if len(dr['class_ids']) == 0: print('No instance is detected.')
def dump(self, output, fixup=False):
"""Dumps WAD to output. Replaces {titleid} and {titleversion} if in filename.
Passing "fixup=True" will repair the common-key index and the certificate chain
"""
if self.ticket.get_titleid().startswith("00030"):
raise Exception("Can't pack DSi Title as WAD.")
output = output.format(titleid=self.tmd.get_titleid(), titleversion=self.tmd.hdr.titleversion)
if fixup:
if self.ticket.hdr.ckeyindex > 2: # Common key index too high
print("Fixing Common key index...")
self.ticket.hdr.ckeyindex = 0
if not self.correct_cert_order: # Fixup certificate chain
print("Fixing Certificate chain...")
ca_cert = None
tmd_cert = None
cetk_cert = None
for cert in self.tmd.certificates + self.ticket.certificates:
if cert.get_name() == "CA00000001" or cert.get_name() == "CA00000002":
ca_cert = cert
if cert.get_name() == "CP00000004" or cert.get_name() == "CP00000007":
tmd_cert = cert
if cert.get_name() == "XS00000003" or cert.get_name() != "XS00000006":
cetk_cert = cert
if not ca_cert:
raise Exception("ERROR: CA Certificate was not found")
if not tmd_cert:
raise Exception("ERROR: CP Certificate was not found")
if not cetk_cert:
raise Exception("ERROR: XS Certificate was not found")
self.certificates = [ca_cert, tmd_cert, cetk_cert]
self.hdr.certchainsize = len(b"".join([x.pack() for x in self.certificates]))
# Header
wad = self.hdr.pack()
wad += utils.align(len(self.hdr))
# Certificate Chain
wad += b"".join([x.pack() for x in self.certificates])
wad += utils.align(self.hdr.certchainsize)
# Ticket
wad += self.ticket.pack()
wad += utils.align(self.hdr.ticketsize)
# TMD
wad += self.tmd.pack()
wad += utils.align(self.hdr.tmdsize)
# Writing WAD
total_content_length = 0
with open(output, "wb") as wad_file:
wad_file.write(wad)
# Not forgetting Contents!
for i, content in enumerate(self.contents):
content_length = 0
for chunk in utils.read_in_chunks(content):
content_length += len(chunk)
wad_file.write(chunk)
wad_file.write(utils.align(content_length))
total_content_length += content_length
# Footer
if self.footer:
wad_file.write(self.footer)
wad_file.write(utils.align(self.hdr.footersize))
def init_uc():
global virtualmemorysize, BASE_ADDR, STACK_ADDR, STACK_SIZE, STACK_START, HOOK_ADDR, mu, startaddr, loaded, apicall_handler
# Calculate required memory
virtualmemorysize = getVirtualMemorySize(sample)
pe = pefile.PE(sample)
BASE_ADDR = pe.OPTIONAL_HEADER.ImageBase # 0x400000
unpacker.BASE_ADDR = BASE_ADDR
STACK_ADDR = 0x0
STACK_SIZE = 1024 * 1024
STACK_START = STACK_ADDR + STACK_SIZE
unpacker.secs += [{
"name": "stack",
"vaddr": STACK_ADDR,
"vsize": STACK_SIZE
}]
HOOK_ADDR = STACK_START + 0x3000 + 0x1000
# Start unicorn emulator with x86-32bit architecture
mu = Uc(UC_ARCH_X86, UC_MODE_32)
if startaddr is None:
startaddr = entrypoint(pe)
loaded = pe.get_memory_mapped_image(ImageBase=BASE_ADDR)
virtualmemorysize = len(loaded)
unpacker.virtualmemorysize = virtualmemorysize
mu.mem_map(BASE_ADDR, align(virtualmemorysize + 0x3000, page_size=4096))
mu.mem_write(BASE_ADDR, loaded)
setup_processinfo(mu)
# Load DLLs
load_dll(mu, "DLLs/KernelBase.dll", 0x73D00000)
load_dll(mu, "DLLs/kernel32.dll", 0x755D0000)
load_dll(mu, "DLLs/ntdll.dll", 0x77400000)
# initialize machine registers
mu.mem_map(STACK_ADDR, STACK_SIZE)
mu.reg_write(UC_X86_REG_ESP, STACK_ADDR + int(STACK_SIZE / 2))
mu.reg_write(UC_X86_REG_EBP, STACK_ADDR + int(STACK_SIZE / 2))
mu.mem_write(mu.reg_read(UC_X86_REG_ESP) + 0x8, bytes([1]))
mu.reg_write(UC_X86_REG_ECX, startaddr)
mu.reg_write(UC_X86_REG_EDX, startaddr)
mu.reg_write(UC_X86_REG_ESI, startaddr)
mu.reg_write(UC_X86_REG_EDI, startaddr)
# init syscall handling and prepare hook memory for return values
apicall_handler = WinApiCalls(BASE_ADDR, virtualmemorysize, HOOK_ADDR,
breakpoints, sample)
mu.mem_map(HOOK_ADDR, 0x1000)
unpacker.secs += [{"name": "hooks", "vaddr": HOOK_ADDR, "vsize": 0x1000}]
hexstr = bytes.fromhex('000000008b0425') + struct.pack(
'<I', HOOK_ADDR) + bytes.fromhex(
'c3'
) # mov eax, [HOOK]; ret -> values of syscall are stored in eax
mu.mem_write(HOOK_ADDR, hexstr)
# handle imports
for lib in pe.DIRECTORY_ENTRY_IMPORT:
for func in lib.imports:
func_name = func.name.decode(
) if func.name is not None else f"no name: 0x{func.address:02x}"
dll_name = lib.dll.decode(
) if lib.dll is not None else "-- unknown --"
imports.add(func_name)
curr_hook_addr = apicall_handler.add_hook(mu, func_name, dll_name)
mu.mem_write(func.address, struct.pack('<I', curr_hook_addr))
# Patch DLLs with hook
apicall_handler.add_hook(mu, "VirtualProtect", "KernelBase.dll",
0x73D00000 + 0x1089f0)
apicall_handler.add_hook(mu, "VirtualAlloc", "KernelBase.dll",
0x73D00000 + 0xd4600)
apicall_handler.add_hook(mu, "VirtualFree", "KernelBase.dll",
0x73D00000 + 0xd4ae0)
apicall_handler.add_hook(mu, "LoadLibraryA", "KernelBase.dll",
0x73D00000 + 0xf20d0)
apicall_handler.add_hook(mu, "GetProcAddress", "KernelBase.dll",
0x73D00000 + 0x102870)
apicall_handler.add_hook(mu, "VirtualProtect", "kernel32.dll",
0x755D0000 + 0x16760)
apicall_handler.add_hook(mu, "VirtualAlloc", "kernel32.dll",
0x755D0000 + 0x166a0)
apicall_handler.add_hook(mu, "VirtualFree", "kernel32.dll",
0x755D0000 + 0x16700)
apicall_handler.add_hook(mu, "LoadLibraryA", "kernel32.dll",
0x755D0000 + 0x157b0)
apicall_handler.add_hook(mu, "GetProcAddress", "kernel32.dll",
0x755D0000 + 0x14ee0)
# Add hooks
mu.hook_add(UC_HOOK_CODE, hook_code)
mu.hook_add(UC_HOOK_MEM_READ | UC_HOOK_MEM_WRITE | UC_HOOK_MEM_FETCH,
hook_mem_access)
mu.hook_add(UC_HOOK_MEM_READ_UNMAPPED | UC_HOOK_MEM_WRITE_UNMAPPED,
hook_mem_invalid)
def setup_processinfo(mu):
global TEB_BASE, PEB_BASE
TEB_BASE = 0x200000
PEB_BASE = TEB_BASE + 0x1000
LDR_PTR = PEB_BASE + 0x1000
LIST_ENTRY_BASE = LDR_PTR + 0x1000
teb = TEB(
-1, # fs:00h
STACK_START, # fs:04h
STACK_START - STACK_SIZE, # fs:08h
0, # fs:0ch
0, # fs:10h
0, # fs:14h
TEB_BASE, # fs:18h (teb base)
0, # fs:1ch
0xdeadbeef, # fs:20h (process id)
0xdeadbeef, # fs:24h (current thread id)
0, # fs:28h
0, # fs:2ch
PEB_BASE, # fs:3ch (peb base)
)
peb = PEB(
0,
0,
0,
0,
0xffffffff,
BASE_ADDR,
LDR_PTR,
)
ntdll_entry = LIST_ENTRY(
LIST_ENTRY_BASE + 12,
LIST_ENTRY_BASE + 24,
0x77400000,
)
kernelbase_entry = LIST_ENTRY(
LIST_ENTRY_BASE + 24,
LIST_ENTRY_BASE + 0,
0x73D00000,
)
kernel32_entry = LIST_ENTRY(
LIST_ENTRY_BASE + 0,
LIST_ENTRY_BASE + 12,
0x755D0000,
)
ldr = PEB_LDR_DATA(
0x30,
0x1,
0x0,
LIST_ENTRY_BASE,
LIST_ENTRY_BASE + 24,
LIST_ENTRY_BASE,
LIST_ENTRY_BASE + 24,
LIST_ENTRY_BASE,
LIST_ENTRY_BASE + 24,
)
teb_payload = bytes(teb)
peb_payload = bytes(peb)
ldr_payload = bytes(ldr)
ntdll_payload = bytes(ntdll_entry)
kernelbase_payload = bytes(kernelbase_entry)
kernel32_payload = bytes(kernel32_entry)
mu.mem_map(TEB_BASE, align(0x5000))
mu.mem_write(TEB_BASE, teb_payload)
mu.mem_write(PEB_BASE, peb_payload)
mu.mem_write(LDR_PTR, ldr_payload)
mu.mem_write(LIST_ENTRY_BASE, ntdll_payload)
mu.mem_write(LIST_ENTRY_BASE + 12, kernelbase_payload)
mu.mem_write(LIST_ENTRY_BASE + 24, kernel32_payload)
mu.windows_tib = TEB_BASE
def asymmetric_jaccard_coef_naive(A, B):
"""
max_diff is trimmed to its minimal box. (max_x, max_y) is the relative position of max_diff to A
:param A:
:param B:
:return: a_j_coef, diff_to_A_x, diff_to_A_y, diff_to_B_x, diff_to_B_y, diff
"""
A_shape_y, A_shape_x = A.shape
B_shape_y, B_shape_x = B.shape
A_sum = A.sum()
if A_shape_x < B_shape_x and A_shape_y < B_shape_y:
A_to_B_coords, a_j_coefs = asymmetric_jaccard_coef_naive_embed(
A, B, A_sum)
elif A_shape_x >= B_shape_x and A_shape_y >= B_shape_y:
B_to_A_coords, a_j_coefs = asymmetric_jaccard_coef_naive_embed(
B, A, A_sum)
A_to_B_coords = -B_to_A_coords
elif A_shape_x < B_shape_x and A_shape_y >= B_shape_y:
A_to_B_coords, a_j_coefs = asymmetric_jaccard_coef_naive_cross(
A, B, A_sum)
elif A_shape_x >= B_shape_x and A_shape_y < B_shape_y:
B_to_A_coords, a_j_coefs = asymmetric_jaccard_coef_naive_cross(
B, A, A_sum)
A_to_B_coords = -B_to_A_coords
else:
raise Exception("Impossible Exception!")
a_j_coef = np.max(a_j_coefs)
coef_argmax = np.where(a_j_coefs == a_j_coef)[0]
if 1 == len(coef_argmax):
ii = coef_argmax[0]
A_to_B_x, A_to_B_y = A_to_B_coords[ii]
else:
A_center_x = (A_shape_x - 1) / 2
A_center_y = (A_shape_y - 1) / 2
B_center_x = (B_shape_x - 1) / 2
B_center_y = (B_shape_y - 1) / 2
closest_center_ii = -1
smallest_center_dist = np.inf
for ii in coef_argmax:
x, y = A_to_B_coords[ii]
center_dist = abs(x + A_center_x -
B_center_x) + abs(y + A_center_y - B_center_y)
if center_dist < smallest_center_dist:
closest_center_ii = ii
smallest_center_dist = center_dist
A_to_B_x, A_to_B_y = A_to_B_coords[closest_center_ii]
A_aligned, B_aligned, aligned_to_B_x, aligned_to_B_y = utils.align(
A, B, A_to_B_x, A_to_B_y)
diff = np.logical_and(B_aligned, np.logical_not(A_aligned))
if diff.any():
diff_y, diff_x = np.where(diff)
diff_x_min = diff_x.min()
diff_x_max = diff_x.max() + 1
diff_y_min = diff_y.min()
diff_y_max = diff_y.max() + 1
diff = diff[diff_y_min:diff_y_max, diff_x_min:diff_x_max]
diff_to_B_x = diff_x_min + aligned_to_B_x
diff_to_B_y = diff_y_min + aligned_to_B_y
diff_to_A_x = diff_to_B_x - A_to_B_x
diff_to_A_y = diff_to_B_y - A_to_B_y
else: # diff is all white, i.e. B is completely covered by A
diff_to_A_x = 0
diff_to_A_y = 0
diff_to_B_x = A_to_B_x
diff_to_B_y = A_to_B_y
diff = np.full_like(A, fill_value=False)
return a_j_coef, diff_to_A_x, diff_to_A_y, diff_to_B_x, diff_to_B_y, utils.erase_noise_point(
diff, 4)
def run_prob_anlg_tran_2x2(prob, anlg, tran):
u1 = prob.matrix[anlg.get("value")[0]]
u2 = prob.matrix[anlg.get("value")[1]]
diff_to_u1_x = None
diff_to_u1_y = None
diff_to_u2_x = None
diff_to_u2_y = None
diff = None
copies_to_u1_x = None
copies_to_u1_y = None
u1_coms_x = None
u1_coms_y = None
u2_coms_x = None
u2_coms_y = None
if "add_diff" == tran.get("name"):
score, diff_to_u1_x, diff_to_u1_y, diff_to_u2_x, diff_to_u2_y, diff = \
asymmetric_jaccard.asymmetric_jaccard_coef(u1, u2)
elif "subtract_diff" == tran.get("name"):
score, diff_to_u2_x, diff_to_u2_y, diff_to_u1_x, diff_to_u1_y, diff = \
asymmetric_jaccard.asymmetric_jaccard_coef(u2, u1)
elif "xor_diff" == tran.get("name"):
score, u1_to_u2_x, u1_to_u2_y = jaccard.jaccard_coef(u1, u2)
score = 1 - score
u1_aligned, u2_aligned, aligned_to_u2_x, aligned_to_u2_y = utils.align(
u1, u2, u1_to_u2_x, u1_to_u2_y)
diff = utils.erase_noise_point(np.logical_xor(u1_aligned, u2_aligned),
4)
diff_to_u2_x = int(aligned_to_u2_x)
diff_to_u2_y = int(aligned_to_u2_y)
diff_to_u1_x = int(diff_to_u2_x - u1_to_u2_x)
diff_to_u1_y = int(diff_to_u2_y - u1_to_u2_y)
elif "upscale_to" == tran.get("name"):
u1_upscaled = transform.upscale_to(u1, u2)
score, _, _ = jaccard.jaccard_coef(u2, u1_upscaled)
elif "duplicate" == tran.get("name"):
scores = []
u1_to_u2_xs = []
u1_to_u2_ys = []
current = u2.copy()
current_to_u2_x = 0
current_to_u2_y = 0
while current.sum():
score_tmp, diff_tmp_to_u1_x, diff_tmp_to_u1_y, diff_tmp_to_current_x, diff_tmp_to_current_y, _ = \
asymmetric_jaccard.asymmetric_jaccard_coef(u1, current)
if score_tmp < 0.6:
break
scores.append(score_tmp)
u1_to_current_x = (-diff_tmp_to_u1_x) - (-diff_tmp_to_current_x)
u1_to_current_y = (-diff_tmp_to_u1_y) - (-diff_tmp_to_current_y)
u1_to_u2_x = u1_to_current_x + current_to_u2_x
u1_to_u2_y = u1_to_current_y + current_to_u2_y
u1_to_u2_xs.append(u1_to_u2_x)
u1_to_u2_ys.append(u1_to_u2_y)
u1_aligned, current_aligned, aligned_to_current_x, aligned_to_current_y = utils.align(
u1, current, u1_to_current_x, u1_to_current_y)
current = utils.erase_noise_point(
np.logical_and(current_aligned, np.logical_not(u1_aligned)), 8)
current_to_u2_x = aligned_to_current_x + current_to_u2_x
current_to_u2_y = aligned_to_current_y + current_to_u2_y
if 1 >= len(scores):
score = 0
copies_to_u1_x = [0]
copies_to_u1_y = [0]
else:
score = min(scores)
copies_to_u1_x = (np.array(u1_to_u2_xs[1:]) -
u1_to_u2_xs[0]).tolist()
copies_to_u1_y = (np.array(u1_to_u2_ys[1:]) -
u1_to_u2_ys[0]).tolist()
elif "rearrange" == tran.get("name"):
score, u1_coms_x, u1_coms_y, u2_coms_x, u2_coms_y = transform.evaluate_rearrange(
u1, u2)
else:
u1_t = transform.apply_unary_transformation(u1, tran)
score, _, _ = jaccard.jaccard_coef(u1_t, u2)
if "mirror" == tran.get("name") or "mirror_rot_180" == tran.get("name"):
# if u1 or u2 is already symmetric, then we shouldn't use mirror tran.
u1_mirror_score, _, _ = jaccard.jaccard_coef(u1_t, u1)
u2_mirror = transform.apply_unary_transformation(u2, tran)
u2_mirror_score, _, _ = jaccard.jaccard_coef(u2_mirror, u2)
if max(u1_mirror_score, u2_mirror_score) > 0.9:
score = 0
prob_anlg_tran_d = assemble_prob_anlg_tran_d(prob,
anlg,
tran,
score,
diff_to_u1_x=diff_to_u1_x,
diff_to_u1_y=diff_to_u1_y,
diff_to_u2_x=diff_to_u2_x,
diff_to_u2_y=diff_to_u2_y,
diff=diff,
copies_to_u1_x=copies_to_u1_x,
copies_to_u1_y=copies_to_u1_y,
u1_coms_x=u1_coms_x,
u1_coms_y=u1_coms_y,
u2_coms_x=u2_coms_x,
u2_coms_y=u2_coms_y)
return prob_anlg_tran_d
def align(self):
self.vertices, self.eigenvectors, self.eigenvalues = align(self.vertices.flatten())
n_neighbors=10,
geod_n_neighbors=10,
embedding_dim=3,
verbose=True).fit()
ptu_time = round(time() - t, 2)
print('ptu time: ', ptu_time)
# Perform Isomap
t = time()
iso = Isomap(n_neighbors=10, n_components=3).fit_transform(exact)
isomap_time = round(time() - t, 2)
print('isomap time: ', isomap_time)
# Align PTU and Isomap to exact parametrization via best isometric
# transformation, and compute errors
ptu = align(ptu, exact)
iso = align(iso, exact)
ptu_error = relative_error(ptu, exact)
iso_error = relative_error(iso, exact)
print('ptu relative error: {}%'.format(ptu_error))
print('isomap relative error: {}%'.format(iso_error))
# Plot results
f = plot_torus(exact,
ptu,
iso,
ptu_time,
isomap_time,
hue='normalized_poinwise_error')
plt.show()
gt['handle_visibility']) == len(gt_class_ids)
print('got handle visibiity.')
else:
result['gt_handle_visibility'] = np.ones_like(
gt_class_ids)
else:
# align gt coord with depth to get RT
if not data in ['coco_val', 'coco_train']:
if len(gt_class_ids) == 0:
print('No gt instance exsits in this image.')
print('\nAligning ground truth...')
start = time.time()
result['gt_RTs'], _, error_message = utils.align(
gt_class_ids, gt_mask, gt_coord, depth, intrinsics,
synset_names, image_path, save_dir + '/' +
'{}_{}_{}_gt_'.format(data, image_path_parsing[-2],
image_path_parsing[-1]))
print(
'New alignment takes {:03f}s.'.format(time.time() -
start))
if len(error_message):
f_log.write(error_message)
result['gt_handle_visibility'] = np.ones_like(gt_class_ids)
## detection
start = time.time()
detect_result = model.detect([image], verbose=0)
r = detect_result[0]
def get_real_original_information(color_path):
# Getting the data set ID and obtaining the corresponding file paths for the
# mask, coordinate map, depth, and meta files
data_id = color_path.name.replace('_color.png', '')
mask_path = color_path.parent / f'{data_id}_mask.png'
depth_path = color_path.parent / f'{data_id}_depth.png'
coord_path = color_path.parent / f'{data_id}_coord.png'
meta_path = color_path.parent / f'{data_id}_meta.txt'
# Loading data
color_image = cv2.imread(str(color_path), cv2.IMREAD_UNCHANGED)
mask_image = cv2.imread(str(mask_path), cv2.IMREAD_UNCHANGED)
coord_map = cv2.imread(str(coord_path), cv2.IMREAD_UNCHANGED)[:, :, :3]
depth_image = cv2.imread(str(depth_path), cv2.IMREAD_UNCHANGED)
# Normalizing data
coord_map = coord_map[:, :, (2, 1, 0)]
# Converting depth to the correct shape and dtype
if len(depth_image.shape) == 3: # encoded depth image
new_depth = np.uint16(depth_image[:, :, 1] * 256) + np.uint16(
depth_image[:, :, 2])
new_depth = new_depth.astype(np.uint16)
depth_image = new_depth
elif len(depth_image.shape) == 2 and depth_image.dtype == 'uint16':
pass # depth is perfecto!
else:
assert False, '[ Error ]: Unsupported depth type'
# flip z axis of coord map
coord_map = np.array(coord_map, dtype=np.float32) / 255
coord_map[:, :, 2] = 1 - coord_map[:, :, 2]
instance_dict = {}
# Loading instance information from meta file
with open(str(meta_path), 'r') as f:
for line in f:
line_info = line.split(' ')
instance_id = int(line_info[0])
class_id = int(line_info[1])
instance_dict[instance_id] = class_id
#print(f"Instance #{instance_id}: {class_id} = {pj.constants.NOCS_CLASSES[class_id]}")
cdata = np.array(mask_image, dtype=np.int32)
# The value of the mask, from 0 to 255, is the class id
instance_ids = list(np.unique(cdata))
instance_ids = sorted(instance_ids)
#print(f'instance_ids: {instance_ids}')
# removing background
assert instance_ids[-1] == 255
del instance_ids[-1]
cdata[cdata == 255] = -1
assert (np.unique(cdata).shape[0] < 20)
num_instance = len(instance_ids)
h, w = cdata.shape
masks = np.zeros([h, w, num_instance], dtype=np.uint8)
coords = np.zeros((h, w, num_instance, 3), dtype=np.float32)
class_ids = np.zeros([num_instance], dtype=np.int_)
scales = np.zeros([num_instance, 3], dtype=np.float32)
# Determing the scales (of the 3d bounding boxes)
with open(str(meta_path), 'r') as f:
lines = f.readlines()
scale_factor = np.zeros((len(lines), 3), dtype=np.float32)
for i, line in enumerate(lines):
words = line[:-1].split(' ')
object_instance_name = words[2]
bbox_file = OBJ_MODEL_DIR / f'{object_instance_name}.txt'
bbox = np.loadtxt(str(bbox_file))
value = bbox #bbox[0, :] - bbox[1, :]
scale_factor[i, :] = value # [a b c] for scale of NOCS [x y z]
scale_factor[i, :] /= np.linalg.norm(scale_factor[i, :])
# Deleting background objects and non-existing objects
instance_dict = {
instance_id: class_id
for instance_id, class_id in instance_dict.items()
if (class_id != 0 and instance_id in instance_ids)
}
i = 0
for instance_id in instance_ids:
if instance_id not in instance_dict.keys():
continue
instance_mask = np.equal(cdata, instance_id)
assert np.sum(instance_mask) > 0
assert instance_dict[instance_id]
masks[:, :, i] = instance_mask
coords[:, :,
i, :] = np.multiply(coord_map,
np.expand_dims(instance_mask, axis=-1))
# class ids are also one-indexed
class_ids[i] = instance_dict[instance_id]
scales[i, :] = scale_factor[instance_id - 1, :]
i += 1
masks = masks[:, :, :i]
coords = coords[:, :, :i, :]
coords = np.clip(coords, 0, 1) # normalize
class_ids = class_ids[:i]
scales = scales[:i]
bboxes = dm.extract_2d_bboxes_from_masks(masks)
scores = [100 for j in range(len(class_ids))]
# RT pickled
RT_json_path = pathlib.Path.cwd() / f'RT_{data_id}.json'
# now obtaining the rotation and translation
if RT_json_path.exists():
data = jt.load_from_json(RT_json_path)
RTs = data['RTs']
else:
# Disable print
disable_print()
RTs, _, _, _ = nocs_utils.align(class_ids, masks, coords, depth_image,
pj.constants.INTRINSICS['REAL'],
pj.constants.REAL_CLASSES, ".", None)
# Enable print
enable_print()
data = {'RTs': RTs}
#jt.save_to_json(RT_json_path, data)
return class_ids, bboxes, masks, coords, RTs, scores, scales, instance_dict
return None
def extract_features(shape, n_bins, n_samples):
'''Extract relevant features, input must be Shape or path'''
if isinstance(shape, Shape):
pass
elif isinstance(shape, Path) or isinstance(shape, str):
shape = Path(shape)
shape = Shape(*read_model(shape))
else:
raise Exception("Input must be Shape or path")
#Make pyvista mesh if not already made
if not shape.pyvista_mesh:
shape.make_pyvista_mesh()
feature_dict = {}
feature_dict["n_vertices"] = shape.n_vertices
feature_dict["n_triangles"] = shape.n_triangles
feature_dict["n_quads"] = shape.n_quads
shape.bounding_rect()
feature_dict["bounding_box"] = shape.bounding_rect_vertices
feature_dict["volume"] = np.maximum(
volume(shape.vertices, shape.element_dict["triangles"]),
0.01) #clamp to avoid 0 volume for 2d models
feature_dict["surface_area"] = shape.pyvista_mesh.area
bounding_box_sides = shape.bounding_rect_vertices.reshape(
(-1, 3)).max(axis=0) - shape.bounding_rect_vertices.reshape(
(-1, 3)).min(axis=0)
bounding_box_sides = np.maximum(
bounding_box_sides,
0.01) #clamp above so no zero division for essentially 2d models
feature_dict["bounding_box_ratio"] = np.max(bounding_box_sides) / np.min(
bounding_box_sides)
feature_dict["compactness"] = np.power(feature_dict["surface_area"], 3) / (
36 * np.pi * np.power(feature_dict["volume"], 2))
feature_dict["bounding_box_volume"] = np.prod(bounding_box_sides)
feature_dict["diameter"] = calculate_diameter(shape.vertices)
*_, eigenvalues = align(shape.vertices)
feature_dict["eccentricity"] = np.max(eigenvalues) / np.maximum(
np.min(eigenvalues), 0.01) #also clamp
#Histograms
feature_dict["angle_three_vertices"] = angle_three_random_vertices(
shape.vertices, n_bins=n_bins, n_samples=n_samples)
feature_dict["barycenter_vertice"] = barycenter_vertice(
shape.vertices,
np.zeros(3, dtype=np.float32),
n_bins=n_bins,
n_samples=n_samples)
feature_dict["two_vertices"] = two_vertices(shape.vertices,
n_bins=n_bins,
n_samples=n_samples)
feature_dict["square_area_triangle"] = square_area_triangle(
shape.vertices, n_bins=n_bins, n_samples=n_samples)
feature_dict["cube_volume_tetrahedron"] = cube_volume_tetrahedron(
shape.vertices, n_bins=n_bins, n_samples=n_samples)
return feature_dict
from parallel_transport_unfolding.ptu import PTU
from utils import align, plot_petals, data_abs_path
# Input manifold embedded in 3D, and exact 2D parametrization
with open(data_abs_path('./data/petals/petals_100D.npy'), 'rb') as f:
X_100D = np.load(f)
with open(data_abs_path('./data/petals/petals_3D.npy'), 'rb') as f:
X_3D = np.load(f)
print('Input pointset shape: ', X_100D.shape)
# Perform PTU
t = time()
ptu = PTU(X=X_100D,
n_neighbors=10,
geod_n_neighbors=10,
embedding_dim=2,
verbose=True).fit()
ptu_time = round(time() - t, 2)
print('ptu time: ', ptu_time)
# Perform Isomap
t = time()
iso = Isomap(n_neighbors=10, n_components=2).fit_transform(X_100D)
isomap_time = round(time() - t, 2)
print('isomap time: ', isomap_time)
iso = align(iso, ptu)
# Plot results
f = plot_petals(X_3D, ptu, iso, ptu_time, isomap_time)
plt.show()
