こんにちは、データエンジニアリングチームの宮崎です。

ユニファではAWSを使用しており、画像やCSVなどの各種データもS3に保存しています。そのため、データ分析したいときはS3のデータをダウンロードして行うのですが、その際ディスクに保存して読み込むのは時間や手間がかかります。そこで、今回はPythonでS3上の画像やCSVデータを直接メモリに展開して読み書きする方法をご紹介します。

PythonからAWSの操作を行うには boto3 というライブラリが提供されています。普段ディスクにファイルをダウンロードやアップロードするときは S3.Client の download_file や upload_file を使用するかと思いますが、メモリにバイトデータをダウンロードやアップロードするときは get_object や put_object を使用します。

画像をPillowで読み書き

Pythonで一般的に使われている画像ライブラリはPillowかと思います。まずはPillowで画像ファイルを読み書きしてみたいと思います。 Pillowでは io.BytesIO を介してバイトデータを読み込むことができます。

import io
import boto3
from PIL import Image

s3_client = boto3.client('s3')

# オブジェクトデータを取得
s3_object = s3_client.get_object(
    Bucket='your_bucket', Key='path/to/elephant.jpg')

# バイトデータを読み込み
image_data = io.BytesIO(s3_object['Body'].read())

# 画像に変換
pil_image = Image.open(image_data)

これで、ディスクを介さずにS3からPillowに画像を読み込めました。続いて、画像をアップロードしてみたいと思います。

# 保存領域を用意
buf = io.BytesIO()

# 用意したメモリに保存
pil_image.save(buf, 'JPEG')

# バイトデータをアップロード
s3_client.put_object(
    Bucket='your_bucket', Key='path/to/elephant_pil.jpg', 
    Body=buf.getvalue())

こちらもアップロードできました。なお、新たにアップロードした elephant_pil.jpg は保存時にJPEGで圧縮されているため、オリジナルとは異なりファイルサイズが小さくなっているのでご注意ください。

response = s3_client.list_objects(
    Bucket='your_bucket', Prefix='path/to')
for content in response['Contents']:
    print(f"{content['Key']}: {content['Size']}")

# path/to/elephant.jpg: 197635
# path/to/elephant_pil.jpg: 96766

画像をTensorFlowで読み書き

続いて、TensorFlowでも画像を読み書きしてみます。 S3上の画像に対して、TensorFlowモデルで推論したい場合などにはPillowよりもTensorFlowで直接読み込んだ方が便利かと思います。 TensorFlowは io.BytesIO を使わずに、直接バイトデータを読み込むことができます。

import tensorflow as tf

# オブジェクトデータを取得
s3_object = s3_client.get_object(
    Bucket='your_bucket', Key='path/to/elephant.jpg')

# 画像に変換
tf_image = tf.io.decode_jpeg(s3_object['Body'].read())

アップロードもダウンロード同様、 io.BytesIO を使わずに直接行えます。

# 画像をバイトデータにエンコード
encoded_tensor = tf.io.encode_jpeg(tf_image, format='rgb')

# バイトデータをアップロード
s3_client.put_object(
    Bucket='your_bucket', Key='path/to/elephant_tf.jpg',
    Body=encoded_tensor.numpy())

Pillowより少ない手間で出来ました。

CSVをPandasで読み書き

画像などのバイナリデータだけでなく、CSVも読み書きできます。次はCSVをPandasに読み込んでDataFrameに変換してみます。

import pandas as pd

# オブジェクトデータを取得
s3_object = s3_client.get_object(
    Bucket='your_bucket', Key='path/to/zoo.csv')

# バイトデータに変換
csv_data = io.BytesIO(s3_object['Body'].read())

# ストリーム位置をリセット
csv_data.seek(0)

# DataFrameに変換
df = pd.read_csv(csv_data, encoding='utf8')

Pillowと同じように io.BytesIO を使って読み込むことができます。一点注意が必要なのは、2回目以降の読み込み時は seek(0) によってストリーム位置を先頭に戻す必要があります。ちなみにPillowは Image.open の中で seek(0) を行なってくれているので不要です。

続いてアップロードも行います。

# 保存領域を準備
buf = io.BytesIO()

# メモリに保存
df.to_csv(buf, index=False)

# バイトデータをアップロード
s3_client.put_object(
    Bucket='your_bucket', Key='path/to/zoo_pd.csv', 
    Body=buf.getvalue())

DataFrameも保存時に小数が丸められるなどして、オリジナルファイルと異なる可能性があるのでご注意ください。

まとめ

今回はPythonでS3上のファイルをメモリから直接読み書きする方法についてご紹介しました。 get_object や put_object を使用してバイトデータをやり取りし、必要に応じて io.BytesIO を介して目的のライブラリに読み込ませるのがコツです。特に大量のファイルを読み書きする際にはディスクを介するよりも効率化されるのではないかと思います。

ユニファで一緒に働く仲間を募集しています！

unifa-e.com