| 469 | |
| 470 | # docs_tag: begin_imagebatchdecoder_nvimagecodec |
| 471 | class ImageBatchDecoder: |
| 472 | def __init__( |
| 473 | self, |
| 474 | input_path, |
| 475 | batch_size, |
| 476 | device_id, |
| 477 | cuda_ctx, |
| 478 | cuda_stream, |
| 479 | cvcuda_perf, |
| 480 | ): |
| 481 | |
| 482 | # docs_tag: begin_init_imagebatchdecoder_nvimagecodec |
| 483 | self.logger = logging.getLogger(__name__) |
| 484 | self.batch_size = batch_size |
| 485 | self.input_path = input_path |
| 486 | self.device_id = device_id |
| 487 | self.total_decoded = 0 |
| 488 | self.batch_idx = 0 |
| 489 | self.cuda_ctx = cuda_ctx |
| 490 | self.cuda_stream = cuda_stream |
| 491 | self.cvcuda_perf = cvcuda_perf |
| 492 | self.decoder = nvimgcodec.Decoder(device_id=device_id) |
| 493 | |
| 494 | # docs_tag: begin_parse_imagebatchdecoder_nvimagecodec |
| 495 | if os.path.isfile(self.input_path): |
| 496 | if os.path.splitext(self.input_path)[1] == ".jpg": |
| 497 | # Read the input image file. |
| 498 | self.file_names = [self.input_path] * self.batch_size |
| 499 | # We will use the nvImageCodec based decoder on the GPU in case of images. |
| 500 | # This will be allocated once during the first run or whenever a batch |
| 501 | # size change happens. |
| 502 | else: |
| 503 | raise ValueError("Unable to read file %s as image." % self.input_path) |
| 504 | |
| 505 | elif os.path.isdir(self.input_path): |
| 506 | # It is a directory. Grab file names of all JPG images. |
| 507 | self.file_names = glob.glob(os.path.join(self.input_path, "*.jpg")) |
| 508 | self.logger.info("Found a total of %d JPEG images." % len(self.file_names)) |
| 509 | |
| 510 | else: |
| 511 | raise ValueError( |
| 512 | "Unknown expression given as input_path: %s." % self.input_path |
| 513 | ) |
| 514 | |
| 515 | # docs_tag: end_parse_imagebatchdecoder_nvimagecodec |
| 516 | |
| 517 | # docs_tag: begin_batch_imagebatchdecoder_nvimagecodec |
| 518 | self.file_name_batches = [ |
| 519 | self.file_names[i : i + self.batch_size] # noqa: E203 |
| 520 | for i in range(0, len(self.file_names), self.batch_size) |
| 521 | ] |
| 522 | # docs_tag: end_batch_imagebatchdecoder_nvimagecodec |
| 523 | |
| 524 | self.max_image_size = 1024 * 1024 * 3 # Maximum possible image size. |
| 525 | |
| 526 | self.logger.info( |
| 527 | "Using nvImageCodec decoder version: %s" % nvimgcodec.__version__ |
| 528 | ) |