Maximum parts in Joint Estimation Model

[Mayankm96]

If I understood the paper correctly, in the joint estimation model, the joint class prediction predicts which joint index a particular point belongs to (with 0 being not a member of a joint). So for a K-links system with K-1 joints, the joint classifier should have an output of shape (K, ) for each point.

When I checked the code for the architecture (Line-129) by default the joint_est_model() takes the argument of maximum parts to be 3 (i.e. K=3) and it is not being set to the current category's number of parts.

This may make sense for the eyeglasses category, but for others like ovens (K=2) and drawers (K=4), it doesn't. On inspecting the provided meta-graphs for these two categories, the output shape is always (3,). I changed the code to account for this but of course, since the graph now is different, the checkpoints won't load.

Assign requires shapes of both tensors to match. lhs shape= [1,128,2] rhs shape= [1,128,3]
	 [[Node: save/Assign_287 = Assign[T=DT_FLOAT, _class=["loc:@SPFN/joint_net/fc4_3/weights"], use_locking=true, validate_shape=true, _device="/job:localhost/replica:0/task:0/device:GPU:0"](SPFN/joint_net/fc4_3/weights, save/RestoreV2/_575)]]
	 [[Node: save/RestoreV2/_582 = _Send[T=DT_FLOAT, client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device_incarnation=1, tensor_name="edge_588_save/RestoreV2", _device="/job:localhost/replica:0/task:0/device:CPU:0"](save/RestoreV2:291)]]
****

Are the provided checkpoints and code correct only for the eyeglasses category while for others it isn't?

Edit: It still works for the oven (K=2) because the last logit is always 0, so when the arg-max is performed, one always get classification as {0, 1}. However, it seems more of an issue for drawers where K=4 where output should be in {0, 1, 2, 3} and not {0, 1, 2}.