FastAPIを使用した画像処理APIの実装

Python

概要

画像を入力とし, 画像を返すAPIをFastAPIで実装する. (FastAPIでAPI部分を実装するより, jsでStream APIを使用するほうが大変であった)

実装

準備

$pip install pillow fastapi jinja2 aiofiles python-multipart uvicorn

API側の実装

画像処理の内容は入力された画像をタイル状に並べるだけのかんたんな処理

UploadFileのリストとして受け取り, StreamingResponseとして返す.

データの読み書きはByteIOを使用して行う.

ByteIOを使用してpillowの画像を返す場合は, 返す前にseekで先頭に戻す必要がある.

@app.post('/api/image-processing')
async def create_image_processing(files: List[UploadFile] = File(...)):
    # open image
    bytes_io = BytesIO(files[0].file.read())
    image = Image.open(bytes_io).convert('RGB')

    # image processing
    data = np.array(image)
    h, w, _ = data.shape
    h = int(h // 2) * 2
    w = int(w // 2) * 2
    data = data[:h, :w, :] \
        .reshape(h // 2, 2, w // 2, 2, -1) \
        .transpose(1, 0, 3, 2, 4) \
        .reshape(h, w, -1)
    content = BytesIO()
    Image.fromarray(data).save(content, format='png')
    content.seek(0)

    # response
    return StreamingResponse(content, media_type='image/png')

フロント側の実装

画像を送信して, そのレスポンスのReadableStreamを処理して表示する処理.

MDN Web Docsのサンプルをそのまま使用している.

postBtn.addEventListener("click", () => {
  //
  if (!imageFile) {
    console.error("no image file");
    return;
  }

  const body = new FormData();
  body.append("files", imageFile);

  fetch("/api/image-processing", {
    method: "POST",
    body: body,
  })
    .then((resp) => {
      const reader = resp.body.getReader();

      return new ReadableStream({
        start(controller) {
          return pump();

          function pump() {
            return reader.read().then(({ done, value }) => {
              if (done) {
                controller.close();
                return;
              }
              controller.enqueue(value);
              return pump();
            });
          }
        },
      });
    })
    .then((stream) => new Response(stream))
    .then((resp) => resp.blob())
    .then((blob) => {
      document.getElementById("output-img").src = URL.createObjectURL(blob);
    });
});

出力結果

f:id:sh1m088io:20210301115128p:plain
画像処理API 例

リポジトリ

github.com

参考資料