SJingjia / PlaneRecTR

PlaneRecTR: Unified Query Learning for 3D Plane Recovery from a Single View
Apache License 2.0
24 stars 1 forks source link

RuntimeError: The size of tensor a (256) must match the size of tensor b (1088) at non-singleton dimension 2 #5

Closed HardikJain02 closed 5 months ago

HardikJain02 commented 5 months ago

Cannot Inference on these pretrained models: PlaneRecTR_swinb_pretrained.pth & PlaneRecTR_hrnet32_pretrained.pth

Traceback (most recent call last):
  File "/workspace/hardik/PlaneRecTR/demo/demo.py", line 211, in <module>
    predictions = demo.run_on_image(img)
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/hardik/PlaneRecTR/demo/predictor.py", line 54, in run_on_image
    predictions = self.predictor(image)
                  ^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/hardik/PlaneRecTR/detectron2/detectron2/engine/defaults.py", line 319, in __call__
    predictions = self.model([inputs])[0]
                  ^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/planerectr/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/planerectr/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/hardik/PlaneRecTR/demo/../PlaneRecTR/PlaneRecTR_model.py", line 299, in forward
    plane_seg, inferred_planes_depth, inferred_seg_depth, valid_param = retry_if_cuda_oom(self.plane_inference)(mask_cls_result, mask_pred_result, param_pred_result, depth_pred_result)
                                                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/hardik/PlaneRecTR/detectron2/detectron2/utils/memory.py", line 70, in wrapped
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/hardik/PlaneRecTR/demo/../PlaneRecTR/PlaneRecTR_model.py", line 367, in plane_inference
    inferred_seg_depth = (segmentation * valid_depth_pred).sum(0) # [h, w]
                          ~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (256) must match the size of tensor b (1088) at non-singleton dimension 2
HardikJain02 commented 5 months ago

Basic resizing worked.

resized_valid_depth_pred = F.interpolate(valid_depth_pred.unsqueeze(0), size=(192, 256), mode='bilinear', align_corners=False).squeeze(0)
            segmentation = (plane_seg[:valid_num].argmax(dim=0)[:,:,None] == torch.arange(valid_num).to(plane_seg)).permute(2, 0, 1) # [h, w, 1] == []  -> [valid_plane_num, h, w]
            print(f"Segmentation size: {segmentation.size()}, Valid_depth_pred: {resized_valid_depth_pred.size()}")
            inferred_seg_depth = (segmentation * resized_valid_depth_pred).sum(0) # [h, w]