95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
import time
|
|
from tbd_dataloader import TBDv2Dataset
|
|
|
|
import numpy as np
|
|
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
|
|
def point2screen(points):
|
|
K = [607.13232421875, 0.0, 638.6468505859375, 0.0, 607.1067504882812, 367.1607360839844, 0.0, 0.0, 1.0]
|
|
K = np.reshape(np.array(K), [3, 3])
|
|
rot_points = np.array(points) + np.array([0, 0.2, 0])
|
|
rot_points = rot_points
|
|
points_camera = rot_points.reshape(3, 1)
|
|
|
|
project_matrix = np.array(K).reshape(3, 3)
|
|
points_prj = project_matrix.dot(points_camera)
|
|
points_prj = points_prj.transpose()
|
|
if not points_prj[:, 2][0] == 0.0:
|
|
points_prj[:, 0] = points_prj[:, 0] / points_prj[:, 2]
|
|
points_prj[:, 1] = points_prj[:, 1] / points_prj[:, 2]
|
|
points_screen = points_prj[:, :2]
|
|
assert points_screen.shape == (1, 2)
|
|
points_screen = points_screen.reshape(-1)
|
|
return points_screen
|
|
|
|
if __name__ == '__main__':
|
|
data = TBDv2Dataset(number_frames_to_sample=1, resize_img=None)
|
|
index = np.random.randint(0, len(data))
|
|
start = time.time()
|
|
(
|
|
kinect_imgs, # <- len x 720 x 1280 x 3
|
|
tracker_imgs,
|
|
battery_imgs,
|
|
skele1,
|
|
skele2,
|
|
bbox,
|
|
tracker_skeID_sample, # <- This is the tracker skeleton ID
|
|
tracker2d,
|
|
label,
|
|
experiment_id, # From here for debugging
|
|
timestep,
|
|
obj_id, # <- This is the object ID as a string
|
|
) = data[index]
|
|
end = time.time()
|
|
print(f"Time for one sample: {end-start}")
|
|
|
|
img = kinect_imgs[-1]
|
|
bbox = bbox[-1]
|
|
print(label.shape)
|
|
|
|
print(skele1.shape)
|
|
print(skele2.shape)
|
|
|
|
skele1 = skele1[-1, :,:]
|
|
skele2 = skele2[-1, :,:]
|
|
|
|
print(skele1.shape)
|
|
|
|
|
|
|
|
# reshape img from c, h, w to h, w, c
|
|
img = img.permute(1, 2, 0)
|
|
|
|
fig, ax = plt.subplots(1)
|
|
ax.imshow(img)
|
|
print(bbox[0], bbox[1], bbox[2], bbox[3]) # t(top left x, top left y, width, height)
|
|
top_left_x, top_left_y, width, height = bbox[0], bbox[1], bbox[2], bbox[3]
|
|
x_min, y_min, x_max, y_max = bbox[0], bbox[1], bbox[2], bbox[3]
|
|
|
|
|
|
|
|
|
|
for i in range(26):
|
|
print(skele1[i,0], skele1[i,1])
|
|
print(skele1[i,:].shape)
|
|
print(point2screen(skele1[i,:]))
|
|
x, y = point2screen(skele1[i,:])[0], point2screen(skele1[i,:])[1]
|
|
ax.text(x, y, f"{i}", fontsize=5, color='w')
|
|
|
|
wedge = patches.Wedge((x,y), 10, 0, 360, width=10, color='b')
|
|
ax.add_patch(wedge)
|
|
|
|
for i in range(26):
|
|
x, y = point2screen(skele2[i,:])[0], point2screen(skele2[i,:])[1]
|
|
ax.text(x, y, f"{i}", fontsize=5, color='w')
|
|
wedge = patches.Wedge((point2screen(skele2[i,:])[0], point2screen(skele2[i,:])[1]), 10, 0, 360, width=10, color='r')
|
|
ax.add_patch(wedge)
|
|
|
|
# Create a Rectangle patch
|
|
# rect = patches.Rectangle((top_left_x, top_left_y-height), width, height, linewidth=1, edgecolor='r', facecolor='none')
|
|
# ax.add_patch(rect)
|
|
# rect = patches.Rectangle((x_min, y_max), x_max-x_min, y_max-y_min, linewidth=1, edgecolor='g', facecolor='none')
|
|
# ax.add_patch(rect)
|
|
fig.savefig(f"bbox_{obj_id}_{index}_{experiment_id}.png")
|