| 154 | |
| 155 | # docs_tag: begin_imp_nvvideodecoder |
| 156 | class nvVideoDecoder: |
| 157 | def __init__(self, enc_file, device_id, cuda_ctx, stream): |
| 158 | """ |
| 159 | Create instance of HW-accelerated video decoder. |
| 160 | :param enc_file: Full path to the MP4 file that needs to be decoded. |
| 161 | :param device_id: id of video card which will be used for decoding & processing. |
| 162 | :param cuda_ctx: A cuda context object. |
| 163 | """ |
| 164 | self.device_id = device_id |
| 165 | self.cuda_ctx = cuda_ctx |
| 166 | self.input_path = enc_file |
| 167 | self.stream = stream |
| 168 | # Demuxer is instantiated only to collect required information about |
| 169 | # certain video file properties. |
| 170 | self.nvDemux = nvvc.PyNvDemuxer(self.input_path) |
| 171 | self.nvDec = nvvc.CreateDecoder( |
| 172 | gpuid=0, |
| 173 | codec=self.nvDemux.GetNvCodecId(), |
| 174 | cudacontext=self.cuda_ctx.handle, |
| 175 | cudastream=self.stream.handle, |
| 176 | enableasyncallocations=False, |
| 177 | ) |
| 178 | |
| 179 | self.w, self.h = self.nvDemux.Width(), self.nvDemux.Height() |
| 180 | self.pixelFormat = self.nvDec.GetPixelFormat() |
| 181 | # In case sample aspect ratio isn't 1:1 we will re-scale the decoded |
| 182 | # frame to maintain uniform 1:1 ratio across the pipeline. |
| 183 | sar = 8.0 / 9.0 |
| 184 | self.fixed_h = self.h |
| 185 | self.fixed_w = int(self.w * sar) |
| 186 | |
| 187 | # frame iterator |
| 188 | def generate_decoded_frames(self): |
| 189 | for packet in self.nvDemux: |
| 190 | for decodedFrame in self.nvDec.Decode(packet): |
| 191 | cvcudaTensor = cvcuda.as_tensor( |
| 192 | cvcuda.as_image(decodedFrame.nvcv_image(), cvcuda.Format.U8) |
| 193 | ) |
| 194 | if cvcudaTensor.layout == "NCHW": |
| 195 | # This will re-format the NCHW tensor to a NHWC tensor which will create |
| 196 | # a copy in the CUDA device decoded frame will go out of scope and the |
| 197 | # backing memory will be available by the decoder. |
| 198 | yield cvcuda.reformat(cvcudaTensor, "NHWC") |
| 199 | else: |
| 200 | raise ValueError("Unexpected tensor layout, NCHW expected.") |
| 201 | |
| 202 | def get_next_frames(self, N): |
| 203 | decoded_frames = list(itertools.islice(self.generate_decoded_frames(), N)) |
| 204 | if len(decoded_frames) == 0: |
| 205 | return None |
| 206 | elif len(decoded_frames) == 1: # this case we dont need stack the tensor |
| 207 | return decoded_frames[0] |
| 208 | else: |
| 209 | # convert from list of tensors to a single tensor (NHWC) |
| 210 | tensorNHWC = cvcuda.stack(decoded_frames) |
| 211 | return tensorNHWC |
| 212 | |
| 213 | |