read_frames.py

import torch
import sys
sys.path.append('/root/VideoProcessingFramework/install/bin')
import PyNvCodec as nvc
import PytorchNvCodec as pnvc

gpuID = 0

nvDec = nvc.PyNvDecoder('dima.mp4', gpuID)
print(nvDec.Width(), nvDec.Height())
to_rgb = nvc.PySurfaceConverter(nvDec.Width(), nvDec.Height(), nvc.PixelFormat.NV12, nvc.PixelFormat.RGB, gpuID)
to_planar = nvc.PySurfaceConverter(nvDec.Width(), nvDec.Height(), nvc.PixelFormat.RGB, nvc.PixelFormat.RGB_PLANAR, gpuID)

i = 0
print('starting')
while True:
    # print(i)
    # Obtain NV12 decoded surface from decoder;
    rawSurface = nvDec.DecodeSingleSurface()
    if (rawSurface.Empty()):
        break
    # print('f')
    # Convert to RGB interleaved;
    rgb_byte = to_rgb.Execute(rawSurface)

    # Convert to RGB planar because that's what to_tensor + normalize are doing;
    rgb_planar = to_planar.Execute(rgb_byte)

    # Create torch tensor from it and reshape because
    # pnvc.makefromDevicePtrUint8 creates just a chunk of CUDA memory
    # and then copies data from plane pointer to allocated chunk;
    surfPlane = rgb_planar.PlanePtr()
    surface_tensor = pnvc.makefromDevicePtrUint8(surfPlane.GpuMem(), surfPlane.Width(), surfPlane.Height(), surfPlane.Pitch(), surfPlane.ElemSize())
    
    # import cv2
    # from PIL import Image
    import imageio
    import numpy as np
    img = surface_tensor.view( (3, nvDec.Height(), nvDec.Width()) ).cpu().detach().numpy().transpose(1,2,0).astype(np.uint8) 

    # img = ( img/ 12.92 ) * (img <= 0.04) + (img > 0.04) * (  np.power( (img + 0.055 ) / (1.0 + 0.055), ( 2.4)) )

    # img = np.power(img, 2.2).astype(np.uint8)
    # print(img)
    # Image.fromarray(img).save(f'res/{i:>05}.jpg')
    imageio.imwrite(f'res/{i:>05}.jpg', img )
    break
    # print(i)
    # break
    # print(i, surface_tensor.dtype, surface_tensor.device)
    # surface_tensor.resize_(3, target_h, target_w)

    # This is optional and depends on what you NN expects to take as input
    # Normalize to range desired by NN. Originally it's 
    # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    # But we scale that to [0;255] input, so we multiply normalization coefficients
    # surface_tensor = surface_tensor.type(dtype=torch.cuda.FloatTensor)
    # mean = torch.tensor([123.675, 116.28, 103.53], dtype=torch.float32, device='cuda')
    # std = torch.tensor([58.395, 57.12, 65.025], dtype=torch.float32, device='cuda')
    i += 1