pytorch / vision

Datasets, Transforms and Models specific to Computer Vision
https://pytorch.org/vision
BSD 3-Clause "New" or "Revised" License
16.21k stars 6.95k forks source link

Adding imshow function to directly display image/video without moving tensor off GPU. #4331

Open PeterL1n opened 3 years ago

PeterL1n commented 3 years ago

🚀 Feature

Looking for a imshow function similar to that of cv2 to directly display tensor without moving it off GPU.

Motivation

The main use case is for displaying video frames for real-time applications (such as webcam demo etc).

If the image is already on the graphics card, it should be a way to directly display it. I currently cannot find other tools available to achieve this and I don't know how to implement it. I think it would be a useful function to include in torchvision.

Alternatives

Moving tensor to CPU every frame, but this is very slow for real-time.

NicolasHug commented 3 years ago

Thanks for the proposal @PeterL1n , but I doubt this will ever be in scope for torchvision: at the moment we don't even offer a basic imshow util for CPU images, let alone videos. Perhaps you could try converting the tensor into whatever cv2 accepts and use the cv2 utils instead?

fmassa commented 3 years ago

Thanks for the feature request @PeterL1n .

I agree with @NicolasHug and we don't have plans on providing such functionality in the future.

As a side note, maybe it might be worth copying a whole batch of predictions from GPU to CPU, or to pre-allocate the output CPU tensor in pinned memory so that you can make faster GPU->CPU copies?