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_nhwc = batch.data
7 # Create an empty list to store filenames
8 filenames = []
9 hwctensor_list = []
10 # Iterate through each image to prepare the filenames
11 for img_idx in range(image_tensors_nhwc.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 hwctensor_list.append(image_tensors_nhwc[img_idx].cuda())
20
21 # Pass the image tensors and filenames to the encoder.
22 self.encoder.write(filenames, hwctensor_list)
23 self.cvcuda_perf.pop_range()