github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/examples/ais-etl/etl_torchvision_io.py (about)

     1  """
     2  ETL to transform images using torchvision.
     3  Communication Type: io://
     4  
     5  Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
     6  """
     7  
     8  import io
     9  import sys
    10  
    11  from torchvision import transforms
    12  from PIL import Image
    13  
    14  from aistore import Client
    15  from aistore.sdk import Bucket
    16  from aistore.sdk.const import ETL_COMM_IO
    17  
    18  client = Client("http://192.168.49.2:8080")
    19  
    20  # cannot apply transforms.PILToTensor() as the expected return type is bytes and not tensor
    21  # if you want to convert it to tensor, return it in "bytes-like" object
    22  
    23  
    24  def apply_image_transforms():
    25      transform = transforms.Compose(
    26          [transforms.Resize(256), transforms.CenterCrop(224), transforms.PILToTensor()]
    27      )
    28      input_bytes = sys.stdin.buffer.read()
    29      sys.stdout.buffer.write(transform(Image.open(io.BytesIO(input_bytes))))
    30  
    31  
    32  deps = ["Pillow", "torchvision"]
    33  
    34  # initialize ETL
    35  client.etl(etl_name="etl_torchvision_io").init_code(
    36      transform=apply_image_transforms,
    37      dependencies=deps,
    38      communication_type=ETL_COMM_IO,
    39  )
    40  
    41  # Transform bucket with given ETL name
    42  job_id = client.bucket("from-bck").transform(
    43      etl_name="etl_torchvision_io", to_bck=Bucket("to-bck"), ext={"jpg": "npy"}
    44  )
    45  client.job(job_id).wait()