Hook up PyTorch U-Net model to video

Question

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

Answer 1

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.