Fast APIで画像をアップロードすると"Object of type Tensor is not JSON serializable"というエラーが発生

Question

実現したいこと

• Fast API上に画像をアップロードし、良品と不良品を判定するアプリを作成したいと考えています。

• しかし、画像をアップロードすると "TypeError: Object of type Tensor is not JSON serializable" というエラーが発生します。

• 解決方法のわかる方がおられましたらご教授いただけますと幸いです。情報不足でしたら、コメントいただければと思います。

前提

vscodeで、model2.pyとmain2.pyのファイルを作成しています。
modelは保存済みのVAEモデルをロードして使用しております。

発生している問題・エラーメッセージ

エラーメッセージ
"TypeError: Object of type Tensor is not JSON serializable"

該当のソースコード(python)

from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from model2 import VAE, z_dim
import torch
import cv2
import numpy as np

app = FastAPI()

# モデルのインスタンスを作成
loaded_model = VAE(z_dim)
# 学習したパラメータを読み込む
loaded_model.load_state_dict(torch.load('vae_model_epoch20.pth', map_location=torch.device('cpu')))

# モデルを評価モードに切り替える（推論時に必要）
loaded_model.eval()

# Define anomaly score threshold (adjust as needed)
anomaly_score_threshold = 35.0

async def process_image(file: UploadFile):
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = torch.tensor(image, dtype=torch.float32) / 255.0
    image = image.permute(2, 0, 1).unsqueeze(0)

    # Encode and reconstruct the image
    with torch.no_grad():
        z, ave, log_dev = loaded_model.encoder(image.view(-1, 302 * 403).to(torch.float32))
        reconstructed_image = loaded_model.decoder(z).view(-1).detach().cpu().numpy()

    # 画像の形状を (302, 403) に変形
    reconstructed_image = reconstructed_image.reshape(-1, 302, 403)

    # Calculate anomaly score
    mse = ((image.view(-1, 302, 403).numpy() - reconstructed_image) ** 2).mean()
    kl_divergence = -0.5 * torch.sum(1 + log_dev - ave**2 - log_dev.exp())
    anomaly_score = mse + kl_divergence

    # Determine if the image is defective
    is_defective = anomaly_score > anomaly_score_threshold

    # Convert torch.Tensor to NumPy array for JSON serialization
    image_np = image.squeeze().permute(1, 2, 0).numpy()
    return JSONResponse(content={"is_defective": is_defective, "anomaly_score": anomaly_score.item(), "image": image_np.tolist()})

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    return await process_image(file)

import torch
import torch.nn as nn
import torch.nn.functional as F

class Encoder(nn.Module):
  def __init__(self, z_dim):
    super().__init__()
    self.lr = nn.Linear(302*403, 300)
    self.lr2 = nn.Linear(300, 100)
    self.lr_ave = nn.Linear(100, z_dim)   #average
    self.lr_dev = nn.Linear(100, z_dim)   #log(sigma^2)
    self.relu = nn.ReLU()

  def forward(self, x):
    x = x.view(-1, 302 * 403)  # 入力を平坦化
    x = self.lr(x)
    x = self.relu(x)
    x = self.lr2(x)
    x = self.relu(x)
    ave = self.lr_ave(x)    #average
    log_dev = self.lr_dev(x)    #log(sigma^2)

    ep = torch.randn_like(ave)   #平均0分散1の正規分布に従い生成されるz_dim次元の乱数
    z = ave + torch.exp(log_dev / 2) * ep   #再パラメータ化トリック
    return z, ave, log_dev    # 潜在変数 z とその統計量（平均 ave と対数分散 log_dev）を得る

class Decoder(nn.Module):
  def __init__(self, z_dim):
    super().__init__()
    self.lr = nn.Linear(z_dim, 100)
    self.lr2 = nn.Linear(100, 300)
    self.lr3 = nn.Linear(300, 302*403)
    self.relu = nn.ReLU()

  def forward(self, z):
    x = self.lr(z)
    x = self.relu(x)
    x = self.lr2(x)
    x = self.relu(x)
    x = self.lr3(x)
    x = torch.sigmoid(x)
    return x.view(-1, 1, 302, 403)  

class VAE(nn.Module):
  def __init__(self, z_dim):      # クラスの初期化。Encoder クラスと Decoder クラスのインスタンスを作成。
    super().__init__()
    self.encoder = Encoder(z_dim)
    self.decoder = Decoder(z_dim)

  def forward(self, x):
    z, ave, log_dev = self.encoder(x)
    x = self.decoder(z)
    return x, z, ave, log_dev

def criterion(predict, target, ave, log_dev):
  bce_loss = F.binary_cross_entropy(predict, target, reduction='sum')
  kl_loss = -0.5 * torch.sum(1 + log_dev - ave**2 - log_dev.exp())
  loss = bce_loss + kl_loss
  return loss

# Create an instance of the VAE model
z_dim = 2
model = VAE(z_dim)

試したこと

return JSONResponse(content={"is_defective": is_defective, "anomaly_score": anomaly_score.item(), "image": image_np.tolist()})

上記コードを以下のように修正しましたが、エラーが解決しませんでした。

return JSONResponse(content={"is_defective": is_defective, "anomaly_score": anomaly_score.item(), "image": image.squeeze().permute(1, 2, 0).numpy().tolist()})

Torch.Tensor型だとJSON変換できないため、NumPy型に変換して検討していますが同じエラーが出続けております。過去の質問事例などを見ても解決できず、質問させていただいた次第です。

補足情報（FW/ツールのバージョンなど）

ここにより詳細な情報を記載してください。

Answer 1

json.dumpsを利用して、カスタムのJSONシリアライザ関数を定義することで解決しました。

def convert_to_json_serializable(obj):
    if isinstance(obj, torch.Tensor):
        return obj.tolist()
    raise TypeError(f'Object of type {obj.__class__.__name__} is not JSON serializable')
-------------------------------------------------------------------------------------------------------------------------------
 # Serialize the response using json.dumps with custom serialization function
    response_content = {
        "is_defective": is_defective,
        "anomaly_score": anomaly_score.item(),
        "image": image_np,
        "reconstructed_image": reconstructed_image_np,
    }
    json_response = JSONResponse(content=json.dumps(response_content, default=convert_to_json_serializable))

    return json_response