GameNet-1 / app.py
MAS-AI-0000's picture
Model inference files and hf setup
dcfcfbb verified
raw
history blame
1.88 kB
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.applications.efficientnet import preprocess_input
from PIL import Image
import numpy as np
import json
import io
# Constants
IMG_SIZE = (300, 300)
MODEL_PATH = "GameNetModel.h5"
LABEL_MAP_PATH = "label_to_index.json"
GENRE_MAP_PATH = "game_genre_map.json"
# Load model & mappings
model = load_model(MODEL_PATH)
with open(LABEL_MAP_PATH) as f:
label_to_index = json.load(f)
index_to_label = {v: k for k, v in label_to_index.items()}
with open(GENRE_MAP_PATH) as f:
genre_map = json.load(f)
# Initialize FastAPI app
app = FastAPI()
# Enable CORS (important for frontend calls)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Response schema
class Prediction(BaseModel):
game: str
genre: str
confidence: float
# Inference route
@app.post("/predict", response_model=Prediction)
async def predict(file: UploadFile = File(...)):
try:
image_bytes = await file.read()
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
img = img.resize(IMG_SIZE)
arr = img_to_array(img)
arr = preprocess_input(arr)
arr = np.expand_dims(arr, axis=0)
preds = model.predict(arr)
class_idx = int(np.argmax(preds))
confidence = float(np.max(preds))
game = index_to_label[class_idx]
genre = genre_map.get(game, "Unknown")
return Prediction(game=game, genre=genre, confidence=confidence)
except Exception as e:
return {"error": str(e)}