0
$\begingroup$

I built a U-Net model in PyTorch that is trained on medical images to detect polyps. The purpose of the model is to do semantic segmentation, so it must predict the location + class of polyps.

Now I want to hook the model up to some videos so they can be inferenced. I can't get this to work, because the input size of each frame is different than the input that the model expects.

The frame I read (with CV2) is of size: [1080, 1920, 3]. The model is of size [64, 3, 7, 7]. I figured 64 is the batch size here, 3 is the channels, but what are the 7, 7? The size of the image I should input? I created a pastebin with the model architecture here: https://pastebin.com/XUV35MbE.

Can someone show me how to input my frame into the model to get a prediction? The code I have now is:

capture = cv2.VideoCapture('data/videos/17.mp4')
success, frame = capture.read()
going = True

model.to('cuda')
model.eval()
while going:
    going, frame = capture.read()
    if not going and frame is None:
        continue
        
    frame = torch.tensor(frame).transpose(0,2).type('torch.cuda.FloatTensor')
    results = model(frame)

Edit: the error I get is: RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 3-dimensional input of size [3, 1920, 1080] instead

$\endgroup$
6
  • $\begingroup$ What error are you getting using the code your provided? In addition, when doing inference in pytorch you should use the torch.no_grad() context manager to prevent the calculation of gradients and speed up the forward pass. $\endgroup$
    – Oxbowerce
    Jul 4 at 11:33
  • $\begingroup$ @Oxbowerce thanks for the comment! The error I get is RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 3-dimensional input of size [3, 1920, 1080] instead $\endgroup$ Jul 4 at 11:45
  • $\begingroup$ I think the issue here is more with the number of dimensions than the image size. If you are feeding in the video frames one by one you need to add the extra fourth dimension denoting the batch size. Since the batch size would simply be one in this case using torch.unsqueeze(frame, 0) should work. $\endgroup$
    – Oxbowerce
    Jul 4 at 11:53
  • $\begingroup$ Thanks @Oxbowerce, that did seem to solve the dimension problem. But now I get the error RuntimeError: Sizes of tensors must match except in dimension 3. Got 135 and 136 (The offending index is 0), do you happen to know why that is? $\endgroup$ Jul 4 at 11:56
  • $\begingroup$ That is likely because of the image size. As you are using a U-Net architecture the encoder size will end up with a tensor of size 135 (1080 / 2 / 2 / 2), while the decoder size will end up with a tensor of size 136 (68 * 2). You can solve this by making sure you can divide the height and width by two n times, with n being the number of times you are downsampling. I would expect a height of 1088 to work, so try padding/rescaling your image to that height. $\endgroup$
    – Oxbowerce
    Jul 4 at 12:03
1
$\begingroup$

The RuntimeError error you're getting is caused by an incorrect number of dimensions in the input. The model expects an input with four dimensions with the first one denoting the batch size. If you are feeding in the video frames one by one the batch size would simply be equal to one, using torch.unsqueeze(frame, 0) should work.

The second RuntimeError has to with the size of the tensors you feeding into the model. As you are using a U-Net architecture the encoder size will end up with a tensor of size 135 (1080 / 2 / 2 / 2), while the decoder size will end up with a tensor of size 136 (68 * 2). When the model then tries to concatenate the two tensors it can't because of different sizes. You can solve this by making sure you can divide the height and width by two n times, with n being the number of times you are downsampling. I would expect a height of 1088 to work, so try padding/rescaling your image to that height.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.