2508181758

main.py

@@ -4,16 +4,14 @@ from typing import List, Dict, Any
 import torch
 import torch.nn as nn
 import torch.optim as optim
-# from torchvision import datasets -> datasetsライブラリを使うため削除
 from torchvision import transforms
-from torch.utils.data import DataLoader, Subset
+from torch.utils.data import DataLoader
 import numpy as np
 import base64
 from io import BytesIO
 from PIL import Image
 import random
-import json
-from datasets import load_dataset # ★★★ Hugging Face datasetsライブラリをインポート
+from datasets import load_dataset
 
 # FastAPIアプリケーションインスタンスを作成
 app = FastAPI()
@@ -105,96 +103,84 @@ class PlayerModel(nn.Module):
             x = layer(x)
         return x
 
-# --- グローバル変数とデータ準備 (★★★ ここを大幅に変更) ---
+# --- グローバル変数とデータ準備 (ステートレス対応) ---
+# これらの変数はサーバー起動時に一度だけ初期化され、リクエスト間で変更されない定数として扱う
 device = torch.device("cpu")
-
-# 1. Hugging Face HubからMNISTデータセットをロード
 mnist_dataset = load_dataset("mnist")
+transform = transforms.Compose([transforms.ToTensor()])
 
-# 2. torchvisionのtransformを定義
-transform = transforms.Compose([
-    transforms.ToTensor(),
-    # transforms.Normalize((0.1307,), (0.3081,)) # 必要に応じて正規化
-])
-
-# 3. データセットにtransformを適用する関数を定義
 def apply_transforms(examples):
-    # PIL Imageのリストをテンソルのリストに変換
     examples['image'] = [transform(image.convert("L")) for image in examples['image']]
     return examples
 
-# 4. データセットにtransformを適用
 mnist_dataset.set_transform(apply_transforms)
-
-# 5. DataLoaderを準備
 train_subset = mnist_dataset['train'].select(range(1000))
 train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
 
-# 6. テスト用の画像リストを作成 (DataLoaderと同じ (image, label) タプルの形式を維持)
 test_images = []
-# メモリ使用量を考慮し、テスト画像は1000個に絞る
 test_subset_for_inference = mnist_dataset['test'].shuffle().select(range(1000))
 for item in test_subset_for_inference:
-    # itemは {'image': tensor, 'label': int} という辞書
-    image_tensor = item['image'].unsqueeze(0)  # バッチ次元 (1) を追加
+    image_tensor = item['image'].unsqueeze(0)
     label_tensor = torch.tensor(item['label'])
     test_images.append((image_tensor, label_tensor))
 
-current_enemy = None
-trained_player_model = None
+# --- バックエンドロジック (ステートレス関数) ---
 
-# --- バックエンドロジック ---
 def get_enemy():
-    global current_enemy
-    image, label = random.choice(test_images)
-    current_enemy = {"image": image, "label": label}
+    """新しい敵の画像(base64)と正解ラベルを返す。サーバー側では状態を保持しない。"""
+    image_tensor, label_tensor = random.choice(test_images)
     
-    img_pil = transforms.ToPILImage()(image.squeeze(0))
+    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}
+    return {
+        "image_b64": "data:image/png;base64," + img_str,
+        "label": label_tensor.item()
+    }
 
-def train_player_model(layer_configs: list):
-    global trained_player_model
+def run_inference(layer_configs: list, enemy_image_b64: str, enemy_label: int):
+    """
+    リクエストごとにモデルを構築・訓練し、与えられた敵データで推論を実行する。
+    サーバー側では状態を一切保持しない。
+    """
+    # 1. モデルをその場で構築し、訓練する
     if not layer_configs:
-        return {"success": False, "message": "モデルが空です。"}
-
+        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エポック学習
-            # ★★★ DataLoaderからのデータ受け取り方を変更
-            for batch_idx, batch in enumerate(train_loader):
+        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(f"Epoch {epoch+1} training finished.")
-
-        trained_player_model = model
-        return {"success": True, "message": "モデルの訓練が完了しました！"}
+        print("On-the-fly training for inference finished.")
     except Exception as e:
-        print(f"Error during training: {e}")
-        return {"success": False, "message": f"訓練中にエラーが発生しました: {e}"}
-
-def run_inference():
-    global trained_player_model, current_enemy
-    if trained_player_model is None:
-        return {"error": "モデルが訓練されていません。"}
+        print(f"Error during on-the-fly training: {e}")
+        return {"error": f"推論中のモデル構築・訓練エラー: {e}"}
 
-    image, label = random.choice(test_images)
-    current_enemy = {"image": image, "label": label}
-    
-    model = trained_player_model
+    # 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):
@@ -206,7 +192,6 @@ def run_inference():
         hooks.append(layer.register_forward_hook(get_hook(name)))
 
     with torch.no_grad():
-        image_tensor = current_enemy["image"].to(device)
         output = model(image_tensor)
         
     for h in hooks: h.remove()
@@ -224,34 +209,38 @@ def run_inference():
                 weights[name + '_w'] = layer.weight.cpu().detach().numpy().tolist()
                 weights[name + '_b'] = layer.bias.cpu().detach().numpy().tolist()
 
-    is_correct = (prediction == current_enemy["label"].item())
-
-    img_pil = transforms.ToPILImage()(image.squeeze(0))
-    buffered = BytesIO()
-    img_pil.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    is_correct = (prediction == enemy_label)
 
+    # 4. 結果をクライアントに返す
     return {
-        "prediction": prediction, "label": current_enemy["label"].item(), "is_correct": is_correct,
+        "prediction": prediction, 
+        "label": enemy_label, 
+        "is_correct": is_correct,
         "confidence": confidence,
-        "image_b64": "data:image/png;base64," + img_str,
+        "image_b64": enemy_image_b64, # 受け取った画像をそのまま返す
         "architecture": [{"type": "Input", "name": "input", "shape": [1, 28, 28]}] + model.architecture_info,
         "outputs": intermediate_outputs,
         "weights": weights
     }
 
-# --- FastAPI Endpoints (変更なし) ---
+# --- FastAPI Endpoints ---
 @app.get("/api/get_enemy")
 async def get_enemy_endpoint():
     return get_enemy()
 
-@app.post("/api/train_player_model")
-async def train_player_model_endpoint(layer_configs: List[Dict[str, Any]] = Body(...)):
-    return train_player_model(layer_configs)
-
 @app.post("/api/run_inference")
-async def run_inference_endpoint():
-    return 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")

web/script.js

@@ -34,6 +34,7 @@ let isBattleInProgress = false; // ★★★ バトルループ中のフラグ
 let draggedItem = null; // { type, layer, index }
 let dragOverIndex = null; // 並び替え先のインデックス
 let wasDroppedSuccessfully = false; // ★★★ このフラグを追加
+let currentEnemy = { image_b64: null, label: null }; // ★★★ クライアント側で敵の状態を保持
 let ENEMY_MAX_HP = 100;
 const PLAYER_MAX_HP = 100;
 
@@ -285,14 +286,19 @@ function updateHpBars() {
 }
 
 async function fetchNewEnemy() {
-    // EelからFetch APIに変更
     const response = await fetch('/api/get_enemy');
-    const enemy = await response.json();
+    const enemyData = await response.json();
     
+    // ★★★ グローバル変数に保存
+    currentEnemy = {
+        image_b64: enemyData.image_b64,
+        label: enemyData.label
+    };
+
     enemyMessage.textContent = '野生のMNISTモンスターが現れた！';
-    enemyImage.src = enemy.image_b64;
+    enemyImage.src = currentEnemy.image_b64;
     enemyImage.classList.remove('hidden');
-    await animateBattleLog('', true); // Clear battle log
+    await animateBattleLog('', true);
 }
 
 // --- D&D Functions ---
@@ -891,7 +897,15 @@ async function handleBattle() {
         await animateBattleLog('新たな敵をスキャン... 推論実行...');
 
         // EelからFetch APIに変更
-        const inferenceResponse = await fetch('/api/run_inference', { method: 'POST' });
+        const inferenceResponse = await fetch('/api/run_inference', {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({
+                layer_configs: playerLayers,
+                enemy_image_b64: currentEnemy.image_b64,
+                enemy_label: currentEnemy.label,
+            }),
+        });
         const result = await inferenceResponse.json();
 
         if (result.error) {
@@ -921,6 +935,7 @@ async function handleBattle() {
 
         if (enemyHP > 0 && playerHP > 0) {
             await sleep(1500);
+            await fetchNewEnemy(); // 次の敵を準備
         }
     }