Image Encoding using nvImageCodec

The image batch encoder is responsible for saving image tensors to the disk as JPG images. The actual encoding is done in batches using the nvImageCodec library. The image encoder is generic enough to be across the sample applications. The code associated with this class can be found in the samples/common/python/nvcodec_utils.py file.

The image batch encoder is a relatively simple class. Here is how its __init__ method is defined.

Once the initialization is complete, we encode the images in the __call__ method. Since the Batch object is passed, we have information of the data, its batch index and the original file name used to read the data.

 1def __call__(self, batch):
 2    self.cvcuda_perf.push_range("encoder.nvimagecodec")
 3
 4    assert isinstance(batch.data, torch.Tensor)
 5
 6    image_tensors_nchw = batch.data
 7    # Create an empty list to store filenames
 8    filenames = []
 9    chwtensor_list = []
10    # Iterate through each image to prepare the filenames
11    for img_idx in range(image_tensors_nchw.shape[0]):
12        img_name = os.path.splitext(os.path.basename(batch.fileinfo[img_idx]))[0]
13        results_path = os.path.join(self.output_path, f"out_{img_name}.jpg")
14        self.logger.info(f"Preparing to save the image to: {results_path}")
15        # Add the filename to the list
16        filenames.append(results_path)
17        # Add the image tensor CAI to a CAI list from an NCHW tensor
18        # (this was a stacked tensor if N images)
19        chwtensor_list.append(image_tensors_nchw[img_idx].cuda())
20
21    # Pass the image tensors and filenames to the encoder.
22    self.encoder.write(filenames, chwtensor_list)
23    self.cvcuda_perf.pop_range()