Tramac / Fast-SCNN-pytorch

A PyTorch Implementation of Fast-SCNN: Fast Semantic Segmentation Network
Apache License 2.0
379 stars 93 forks source link

Testing on a video #40

Closed SamSamhuns closed 3 years ago

SamSamhuns commented 3 years ago

Is there a possible implementation for running the segmentation network on a video through open cv? Any help would be much appreciated.

When I try to load a video through open cv and use my webcam as the source, I get the following error:

File "/Users/user/Downloads/Fast-SCNN-pytorch/models/fast_scnn.py", line 213, in forward out = higher_res_feature + lower_res_feature RuntimeError: The size of tensor a (90) must match the size of tensor b (92) at non-singleton dimension 2

The edited code in demo.py is:

def demo():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # output folder
    if not os.path.exists(args.outdir):
        os.makedirs(args.outdir)

    # image transform
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])

    model = get_fast_scnn(args.dataset, pretrained=True,
                          root=args.weights_folder, map_cpu=args.cpu).to(device)
    print('Finished loading model!')
    model.eval()

    capture = cv2.VideoCapture(0)
    while capture.isOpened():
        ret, frame = capture.read()
        if ret:
            frame = Image.fromarray(frame.astype('uint8')).convert('RGB')
            image = transform(frame).unsqueeze(0).to(device)

            with torch.no_grad():
                result = model(image)
            pred = torch.argmax(result[0], 1).squeeze(0).cpu().data.numpy()
            pred = pred.astype('float64')
            cv2.imshow("frame", pred)

            key = cv2.waitKey(1)
            if key & 0xFF == ord('q'):
                break
        else:
            break
Tramac commented 3 years ago

In theory, it is possible. It seems you need to resize the image.

SamSamhuns commented 3 years ago

Thank you, my webcam's default resolution was 1200x720. Resized the images to 900x900 and the forward pass works now.