main.py

from fastapi import FastAPI, Body
from fastapi.staticfiles import StaticFiles
from typing import List, Dict, Any
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np
import base64
from io import BytesIO
from PIL import Image
import random
from datasets import load_dataset

# FastAPIアプリケーションインスタンスを作成
app = FastAPI()

# --- 動的なプレイヤーモデル (変更なし) ---
class PlayerModel(nn.Module):
    def __init__(self, layer_configs):
        super(PlayerModel, self).__init__()
        self.layers = nn.ModuleList()
        self.architecture_info = []
        self.hookable_layers = {}

        in_channels = 1
        feature_map_size = 28
        is_flattened = False
        
        for i, config in enumerate(layer_configs):
            layer_type = config['type']
            name = f"{layer_type.lower()}_{len([info for info in self.architecture_info if info['type'] == layer_type])}"

            if layer_type in ['Conv2d', 'MaxPool2d', 'AvgPool2d']:
                is_flattened = False
                if layer_type == 'Conv2d':
                    out_channels = config['params']['out_channels']
                    kernel_size = config['params']['kernel_size']
                    layer = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
                    self.layers.append(layer)
                    self.hookable_layers[name] = layer
                    in_channels = out_channels
                    self.architecture_info.append({"type": "Conv2d", "name": name, "shape": [out_channels, feature_map_size, feature_map_size]})
                else:
                    kernel_size = config['params']['kernel_size']
                    if layer_type == 'MaxPool2d':
                        layer = nn.MaxPool2d(kernel_size=kernel_size, stride=kernel_size)
                    else:
                        layer = nn.AvgPool2d(kernel_size=kernel_size, stride=kernel_size)
                    self.layers.append(layer)
                    self.hookable_layers[name] = layer
                    feature_map_size //= kernel_size
                    self.architecture_info.append({"type": layer_type, "name": name, "shape": [in_channels, feature_map_size, feature_map_size]})
            elif layer_type in ['ReLU', 'Dropout']:
                if layer_type == 'ReLU':
                    self.layers.append(nn.ReLU())
                else:
                    p = config['params']['p']
                    self.layers.append(nn.Dropout(p=p))
                self.architecture_info.append({"type": layer_type, "name": name})
            elif layer_type == 'Flatten':
                if not is_flattened:
                    layer = nn.Flatten()
                    self.layers.append(layer)
                    self.hookable_layers[name] = layer
                    flat_features = in_channels * feature_map_size * feature_map_size
                    in_channels = flat_features
                    self.architecture_info.append({"type": "Flatten", "name": name, "shape": [flat_features]})
                    is_flattened = True
            elif layer_type in ['Linear', 'ResidualBlock']:
                if not is_flattened:
                    auto_flatten_name = f"auto_flatten_{i}"
                    self.layers.append(nn.Flatten())
                    flat_features = in_channels * feature_map_size * feature_map_size
                    in_channels = flat_features
                    self.architecture_info.append({"type": "Flatten", "name": auto_flatten_name, "shape": [flat_features]})
                    is_flattened = True
                if layer_type == 'Linear':
                    out_features = config['params']['out_features']
                    layer = nn.Linear(in_channels, out_features)
                    in_channels = out_features
                else:
                    features = in_channels
                    layer = nn.Linear(features, features)
                self.layers.append(layer)
                self.hookable_layers[name] = layer
                self.architecture_info.append({"type": layer_type, "name": name, "shape": [in_channels]})

        if not self.layers or not isinstance(self.layers[-1], nn.Linear) or self.layers[-1].out_features != 10:
            if not is_flattened:
                 self.layers.append(nn.Flatten())
                 final_in_features = in_channels * feature_map_size * feature_map_size
            else:
                 final_in_features = in_channels
            output_layer = nn.Linear(final_in_features, 10)
            self.layers.append(output_layer)
            self.hookable_layers["linear_output"] = output_layer
            self.architecture_info.append({"type": "Linear", "name": "linear_output", "shape": [10]})

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

# --- グローバル変数とデータ準備 (ステートレス対応) ---
# これらの変数はサーバー起動時に一度だけ初期化され、リクエスト間で変更されない定数として扱う
device = torch.device("cpu")
mnist_dataset = load_dataset("mnist")
transform = transforms.Compose([transforms.ToTensor()])

def apply_transforms(examples):
    examples['image'] = [transform(image.convert("L")) for image in examples['image']]
    return examples

mnist_dataset.set_transform(apply_transforms)
train_subset = mnist_dataset['train'].select(range(1000))
train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)

test_images = []
test_subset_for_inference = mnist_dataset['test'].shuffle().select(range(1000))
for item in test_subset_for_inference:
    image_tensor = item['image'].unsqueeze(0)
    label_tensor = torch.tensor(item['label'])
    test_images.append((image_tensor, label_tensor))

# --- バックエンドロジック (ステートレス関数) ---

def get_enemy():
    """新しい敵の画像(base64)と正解ラベルを返す。サーバー側では状態を保持しない。"""
    image_tensor, label_tensor = random.choice(test_images)
    
    img_pil = transforms.ToPILImage()(image_tensor.squeeze(0))
    buffered = BytesIO()
    img_pil.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    
    return {
        "image_b64": "data:image/png;base64," + img_str,
        "label": label_tensor.item()
    }

def run_inference(layer_configs: list, enemy_image_b64: str, enemy_label: int):
    """
    リクエストごとにモデルを構築・訓練し、与えられた敵データで推論を実行する。
    サーバー側では状態を一切保持しない。
    """
    # 1. モデルをその場で構築し、訓練する
    if not layer_configs:
        return {"error": "モデルが空です。"}
    try:
        model = PlayerModel(layer_configs).to(device)
        optimizer = optim.Adam(model.parameters(), lr=0.001)
        loss_fn = nn.CrossEntropyLoss()
        
        model.train()
        for epoch in range(3): # 毎回3エポック学習
            for batch in train_loader:
                data, target = batch['image'].to(device), batch['label'].to(device)
                optimizer.zero_grad()
                output = model(data)
                loss = loss_fn(output, target)
                loss.backward()
                optimizer.step()
        print("On-the-fly training for inference finished.")
    except Exception as e:
        print(f"Error during on-the-fly training: {e}")
        return {"error": f"推論中のモデル構築・訓練エラー: {e}"}

    # 2. クライアントから送られてきた敵画像で推論する
    model.eval()

    # Base64文字列から画像テンソルにデコード
    try:
        header, encoded = enemy_image_b64.split(",", 1)
        image_data = base64.b64decode(encoded)
        image_pil = Image.open(BytesIO(image_data)).convert("L")
        image_tensor = transforms.ToTensor()(image_pil).unsqueeze(0).to(device)
    except Exception as e:
        print(f"Error decoding enemy image: {e}")
        return {"error": f"敵画像のデコードエラー: {e}"}

    # 3. 推論と中間出力のキャプチャ
    intermediate_outputs = {}
    hooks = []
    def get_hook(name):
        def hook(model, input, output):
            intermediate_outputs[name] = output.detach().cpu().clone().numpy().tolist()
        return hook

    for name, layer in model.hookable_layers.items():
        hooks.append(layer.register_forward_hook(get_hook(name)))

    with torch.no_grad():
        output = model(image_tensor)
        
    for h in hooks: h.remove()
    
    probabilities = torch.nn.functional.softmax(output, dim=1)
    prediction = torch.argmax(probabilities, dim=1).item()
    confidence = probabilities[0, prediction].item()
    
    intermediate_outputs['input'] = image_tensor.cpu().numpy().tolist()
    
    weights = {}
    for name, layer in model.hookable_layers.items():
        if isinstance(layer, (nn.Linear, nn.Conv2d)):
            if hasattr(layer, 'weight') and hasattr(layer, 'bias'):
                weights[name + '_w'] = layer.weight.cpu().detach().numpy().tolist()
                weights[name + '_b'] = layer.bias.cpu().detach().numpy().tolist()

    is_correct = (prediction == enemy_label)

    # 4. 結果をクライアントに返す
    return {
        "prediction": prediction, 
        "label": enemy_label, 
        "is_correct": is_correct,
        "confidence": confidence,
        "image_b64": enemy_image_b64, # 受け取った画像をそのまま返す
        "architecture": [{"type": "Input", "name": "input", "shape": [1, 28, 28]}] + model.architecture_info,
        "outputs": intermediate_outputs,
        "weights": weights
    }

# --- FastAPI Endpoints ---
@app.get("/api/get_enemy")
async def get_enemy_endpoint():
    return get_enemy()

@app.post("/api/run_inference")
async def run_inference_endpoint(payload: Dict[str, Any] = Body(...)):
    """
    クライアントからモデル構成と敵データを受け取り、推論結果を返すエンドポイント。
    """
    layer_configs = payload.get("layer_configs")
    enemy_image_b64 = payload.get("enemy_image_b64")
    enemy_label = payload.get("enemy_label")

    if not all([layer_configs, enemy_image_b64, enemy_label is not None]):
        return {"error": "リクエストのパラメータが不足しています。"}

    return run_inference(layer_configs, enemy_image_b64, enemy_label)

# --- 静的ファイルの配信 ---
app.mount("/", StaticFiles(directory="web", html=True), name="static")