WebアプリからTorchServeに投げた画像を記録する

画像認識系のWebアプリで、実際に処理された画像を記録しておきたい事はあるんじゃないかと思います。その方法を考えました。サーバーはgo言語、クライアント側はreact、画像認識部分はTorchServeを使用しているとします。

まずクライアント側ですが、torchserveはBaseHandlerを継承している場合、画像データをそのまま投げる必要があります。axiosを使った場合、こんな感じでOKです。

index.tsx

import * as React from 'react'
import axios, { AxiosResponse } from 'axios'
import './style.css'
import * as ReactDOM from 'react-dom'

export const Default = () => {
  const canvasRef = React.useRef<HTMLCanvasElement>(null)
  const imgRef = React.useRef<HTMLImageElement>(null)
  const predict = () => {
    if (canvasRef.current !== null && imgRef.current !== null) {
      const cxt = canvasRef.current.getContext('2d')
      if (cxt !== null) {
        cxt.drawImage(imgRef.current, 0, 0)
        canvasRef.current.toBlob((blob) => {
          axios
            .request({
              url: '/predictions/model_name',
              method: 'post',
              headers: {
                'Content-type': 'image/png',
              },
              data: blob,
            })
            .then((x: AxiosResponse) => {
              console.log(JSON.stringify(x.data, null, ' '))
            })
            .catch(console.log)
        })
      }
    }
  }

  return (
    <div>
      <canvas width="1008px" height="756px" ref={canvasRef} />
      <img src="/images/test.jpg" ref={imgRef} style={{ display: 'none' }} />
      <button type="button" onClick={predict}>
        PREDICT
      </button>
    </div>
  )
}

ReactDOM.render(
  <React.StrictMode>
    <Default />
  </React.StrictMode>,
  document.getElementById('root')
)

次にサーバー側ですが、TorchServeのREST APIのURLがhttp://127.0.0.1:8080とすると、次のようにリバースプロキシでリクエストを見て、画像ファイルとして保存します。r.Body = ioutil.NopCloser(io.TeeReader(r.Body, f))の1行がポイントで、リクエストをTorchServeに渡しつつ、リクエストのボディをファイルとして保存しています。

main.go

package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"
	"time"
)

func main() {
	target, err := url.Parse("http://127.0.0.1:8080")
	if err != nil {
		log.Fatal(err)
	}
	proxy := httputil.NewSingleHostReverseProxy(target)
	http.HandleFunc("/predictions/", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
			return
        }
        //ここから
		imagePath := path.Join("/path/to/image_dir", fmt.Sprintf("%s.png", time.Now().Format("20060102150405")))
		f, err := os.Create(imagePath)
		if err == nil {
			defer f.Close()
			r.Body = ioutil.NopCloser(io.TeeReader(r.Body, f))
		} else {
            log.Fatalln(err)
        }
        //ここまで
		proxy.ServeHTTP(w, r)
	})
	http.Handle("/", http.FileServer(http.Dir("./build")))
	http.ListenAndServe(":8000", nil)
}

ファイルを作成してr.Bodyの中身をすり替える処理の部分(コード中のコメントの"ここから"と"ここまで"の範囲)を次のように変えると、Google Cloud Storageへの保存に切り替えることもできます(imagePathがpath.Joinではなくstrings.Joinなのは、Windowsのときセパレータがおかしくなるのを防ぐため)。

ctx := context.Background()
client, err := storage.NewClient(ctx)
if err == nil {
    imagePath := strings.Join([]string{"/path/to/image_dir/", fmt.Sprintf("%s.png", time.Now().Format("20060102150405"))}, "/")
    writer := client.Bucket("backet_name").Object(imagePath).NewWriter(ctx)
    defer writer.Close()
    writer.ContentType = r.Header.Get("Content-Type")
    r.Body = ioutil.NopCloser(io.TeeReader(r.Body, writer))
}