marvinandClaude Sonnet 5
121d8a8ab4
- API: read the upload size limit from MAX_UPLOAD_MEGABYTES (default 15) and keep MAX_UPLOAD_BYTES as the byte value used by /predict - Caddy: read the ACME email and request body limit from CADDY_EMAIL and CADDY_MAX_BODY (default 16MB, ~1 MB above the API limit) - compose: pass the new variables to caddy_proxy and cnn_api Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
430 lines
14 KiB
Python
430 lines
14 KiB
Python
import asyncio
|
|
from collections import defaultdict, deque
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import asynccontextmanager, closing
|
|
from datetime import datetime, timezone
|
|
from enum import Enum
|
|
import io
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import secrets
|
|
import sqlite3
|
|
import time
|
|
|
|
from dotenv import load_dotenv
|
|
from fastapi import Depends, FastAPI, File, Header, HTTPException, Query, Request, UploadFile, WebSocket, WebSocketDisconnect
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from PIL import Image
|
|
from pydantic import BaseModel, Field
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torchvision import transforms
|
|
|
|
from yolo_model_production import convert_prediction, Yolo_model
|
|
from bird_cnn_production import Bird_CNN
|
|
|
|
logger = logging.getLogger("server")
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent
|
|
BUILD_PATH = Path(os.getenv("BUILD_PATH", BASE_DIR / "build_models"))
|
|
DATABASE_DIR = Path(os.getenv("DATABASE_DIR", BASE_DIR / "database"))
|
|
DATABASE = DATABASE_DIR / "reviews.db"
|
|
|
|
load_dotenv(DATABASE_DIR / ".env")
|
|
API_KEY = os.getenv("API_KEY") # if missing, protected endpoints reject every request
|
|
|
|
IMAGE_SIZE_CNN = 64
|
|
IMAGE_SIZE_YOLO = 64
|
|
|
|
# ---- Resource limits (the server is weak, keep everything small and bounded) ----
|
|
# Number of predictions (bird or face) that may run at the same time.
|
|
PARALLEL_INFERENCES = max(1, int(os.getenv("PARALLEL_INFERENCES", "1")))
|
|
# Extra requests allowed to wait for a free slot; anything beyond is rejected.
|
|
QUEUED_INFERENCES = max(0, int(os.getenv("QUEUED_INFERENCES", "1")))
|
|
# PyTorch threads used by EACH running prediction. Total CPU use is roughly
|
|
# PARALLEL_INFERENCES * TORCH_THREADS, so keep the product <= your CPU cores.
|
|
TORCH_THREADS = max(1, int(os.getenv("TORCH_THREADS", "1")))
|
|
MAX_UPLOAD_BYTES = int(max(1, float(os.getenv("MAX_UPLOAD_MEGABYTES", "15"))) * 1024 * 1024)
|
|
MAX_FRAME_BYTES = 1 * 1024 * 1024
|
|
MAX_IMAGE_PIXELS = 20_000_000
|
|
MAX_WEBSOCKETS = 3
|
|
MAX_WEBSOCKETS_PER_IP = 1
|
|
FACE_MAX_FPS = float(os.getenv("FACE_MAX_FPS", "5.5")) # per websocket connection
|
|
FACE_MIN_INTERVAL = 1 / FACE_MAX_FPS
|
|
DECODE_DRAFT_SIZE = (256, 256) # JPEG decodes at reduced scale, still larger than the 64px model input
|
|
|
|
Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
|
|
|
|
origins = [
|
|
"https://marvinkrausser.com",
|
|
]
|
|
|
|
|
|
class BirdSpecies(Enum):
|
|
Common_Kingfisher = 0
|
|
Common_Myna = 1
|
|
House_Crow = 2
|
|
Indian_Peacock = 3
|
|
Indian_Pitta = 4
|
|
Ruddy_Shelduck = 5
|
|
Sarus_Crane = 6
|
|
|
|
|
|
transform_bird = transforms.Compose([
|
|
transforms.Resize(IMAGE_SIZE_CNN),
|
|
transforms.CenterCrop(IMAGE_SIZE_CNN),
|
|
transforms.ToTensor()
|
|
])
|
|
|
|
transform_face = transforms.Compose([
|
|
transforms.Resize((IMAGE_SIZE_YOLO, IMAGE_SIZE_YOLO)),
|
|
transforms.ToTensor()
|
|
])
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
models = {}
|
|
|
|
|
|
def load_models():
|
|
model_bird = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
|
|
model_bird.load_state_dict(torch.load(BUILD_PATH / "bird_cnn", map_location=device, weights_only=True))
|
|
model_bird.to(device).eval()
|
|
|
|
model_face = Yolo_model(c_in=3, boxes=1, grid=6, labels=1, c_hidden=16)
|
|
model_face.load_state_dict(torch.load(BUILD_PATH / "face_detection_yolo", map_location=device, weights_only=True))
|
|
model_face.to(device).eval()
|
|
|
|
models["bird"] = model_bird
|
|
models["face"] = model_face
|
|
|
|
|
|
def init_db():
|
|
DATABASE_DIR.mkdir(parents=True, exist_ok=True)
|
|
with closing(sqlite3.connect(DATABASE)) as conn:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS reviews (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
website TEXT NOT NULL,
|
|
rating INTEGER NOT NULL,
|
|
text TEXT NOT NULL,
|
|
created_at TEXT NOT NULL
|
|
)
|
|
""")
|
|
conn.commit()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
logging.basicConfig(level=logging.INFO)
|
|
if not API_KEY:
|
|
logger.warning("API_KEY is not set: GET/DELETE /review are disabled")
|
|
load_models()
|
|
init_db()
|
|
yield
|
|
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins,
|
|
allow_methods=["GET", "POST", "DELETE"],
|
|
allow_headers=["Authorization", "Content-Type"],
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Overload protection
|
|
# --------------------------------------------------------------------------
|
|
|
|
def client_ip(conn: Request | WebSocket) -> str:
|
|
# Caddy is the only proxy in front of this container and appends the real
|
|
# client address as the last X-Forwarded-For entry.
|
|
forwarded = conn.headers.get("x-forwarded-for")
|
|
if forwarded:
|
|
return forwarded.split(",")[-1].strip()
|
|
return conn.client.host if conn.client else "unknown"
|
|
|
|
|
|
class RateLimiter:
|
|
"""Sliding-window limiter per key (in memory, single process)."""
|
|
|
|
def __init__(self, limit: int, window: float):
|
|
self.limit = limit
|
|
self.window = window
|
|
self.hits = defaultdict(deque)
|
|
|
|
def allow(self, key: str) -> bool:
|
|
now = time.monotonic()
|
|
if len(self.hits) > 10_000:
|
|
self.hits = defaultdict(deque, {
|
|
k: v for k, v in self.hits.items() if v and now - v[-1] < self.window
|
|
})
|
|
q = self.hits[key]
|
|
while q and now - q[0] >= self.window:
|
|
q.popleft()
|
|
if len(q) >= self.limit:
|
|
return False
|
|
q.append(now)
|
|
return True
|
|
|
|
|
|
predict_limiter = RateLimiter(
|
|
limit=max(1, int(os.getenv("PREDICT_RATE_LIMIT", "10"))),
|
|
window=max(1.0, float(os.getenv("PREDICT_RATE_WINDOW", "60"))),
|
|
)
|
|
review_limiter = RateLimiter(
|
|
limit=max(1, int(os.getenv("REVIEW_RATE_LIMIT", "5"))),
|
|
window=max(1.0, float(os.getenv("REVIEW_RATE_WINDOW", "60"))),
|
|
)
|
|
|
|
|
|
def rate_limit(limiter: RateLimiter):
|
|
def dependency(request: Request):
|
|
if not limiter.allow(client_ip(request)):
|
|
raise HTTPException(status_code=429, detail="Too many requests", headers={"Retry-After": "60"})
|
|
return dependency
|
|
|
|
|
|
class Busy(Exception):
|
|
pass
|
|
|
|
|
|
def init_inference_thread():
|
|
# Must run inside each worker thread: with OpenMP the thread count is a
|
|
# per-thread setting, so setting it once in the main thread is not enough.
|
|
torch.set_num_threads(TORCH_THREADS)
|
|
|
|
|
|
class InferenceGate:
|
|
"""Runs predictions on `parallel` worker threads. At most
|
|
`parallel + queued` jobs (running + waiting) are admitted; everything
|
|
else is rejected at once instead of queueing up and eating memory.
|
|
|
|
The models are in eval mode under torch.inference_mode(), so several
|
|
threads can safely run forward passes on the same model object."""
|
|
|
|
def __init__(self, parallel: int, queued: int):
|
|
self.max_pending = parallel + queued
|
|
self.pending = 0
|
|
self.executor = ThreadPoolExecutor(
|
|
max_workers=parallel,
|
|
thread_name_prefix="inference",
|
|
initializer=init_inference_thread,
|
|
)
|
|
|
|
async def run(self, fn, *args):
|
|
if self.pending >= self.max_pending:
|
|
raise Busy
|
|
self.pending += 1
|
|
try:
|
|
return await asyncio.get_running_loop().run_in_executor(self.executor, fn, *args)
|
|
finally:
|
|
self.pending -= 1
|
|
|
|
|
|
gate = InferenceGate(PARALLEL_INFERENCES, QUEUED_INFERENCES)
|
|
|
|
|
|
class InvalidImage(Exception):
|
|
pass
|
|
|
|
|
|
def open_image(data: bytes) -> tuple[Image.Image, tuple[int, int]]:
|
|
"""Open lazily, reject oversized images before decoding, decode JPEGs at
|
|
reduced scale. Returns the image and its ORIGINAL (width, height)."""
|
|
try:
|
|
image = Image.open(io.BytesIO(data))
|
|
original_size = image.size
|
|
if original_size[0] * original_size[1] > MAX_IMAGE_PIXELS:
|
|
raise InvalidImage("Image too large")
|
|
image.draft("RGB", DECODE_DRAFT_SIZE)
|
|
return image.convert("RGB"), original_size
|
|
except InvalidImage:
|
|
raise
|
|
except Exception as e: # PIL raises many different types for bad data
|
|
raise InvalidImage("Not a valid image") from e
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Bird classification
|
|
# --------------------------------------------------------------------------
|
|
|
|
def classify_bird(data: bytes) -> dict:
|
|
image, _ = open_image(data)
|
|
tensor = transform_bird(image).unsqueeze(0).to(device)
|
|
with torch.inference_mode():
|
|
probs = F.softmax(models["bird"](tensor), dim=1)
|
|
confidence, cls = torch.max(probs, dim=1)
|
|
return {
|
|
"class": BirdSpecies(cls.item()).name,
|
|
"confidence": confidence.item()
|
|
}
|
|
|
|
|
|
async def read_limited(file: UploadFile, limit: int) -> bytes:
|
|
data = bytearray()
|
|
while chunk := await file.read(64 * 1024):
|
|
data += chunk
|
|
if len(data) > limit:
|
|
raise HTTPException(status_code=413, detail="File too large")
|
|
return bytes(data)
|
|
|
|
|
|
@app.post("/predict", dependencies=[Depends(rate_limit(predict_limiter))])
|
|
async def predict(file: UploadFile = File(...)):
|
|
data = await read_limited(file, MAX_UPLOAD_BYTES)
|
|
try:
|
|
return await gate.run(classify_bird, data)
|
|
except Busy:
|
|
raise HTTPException(status_code=503, detail="Server busy, try again shortly", headers={"Retry-After": "5"})
|
|
except InvalidImage as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Face detection (websocket)
|
|
# --------------------------------------------------------------------------
|
|
|
|
def detect_faces(data: bytes) -> list:
|
|
image, (W, H) = open_image(data)
|
|
scale_w = W / IMAGE_SIZE_YOLO
|
|
scale_h = H / IMAGE_SIZE_YOLO
|
|
|
|
tensor = transform_face(image).unsqueeze(0).to(device)
|
|
with torch.inference_mode():
|
|
pred = models["face"](tensor)
|
|
bboxes, _, _ = convert_prediction(pred.squeeze(0), tensor.squeeze(0), threshold=0.9)
|
|
|
|
return [
|
|
[int(b[0] * scale_w), int(b[1] * scale_h), int(b[2] * scale_w), int(b[3] * scale_h)]
|
|
for b in bboxes
|
|
]
|
|
|
|
|
|
active_websockets = defaultdict(int)
|
|
|
|
|
|
@app.websocket("/predict_face")
|
|
async def predict_face(websocket: WebSocket):
|
|
ip = client_ip(websocket)
|
|
|
|
# CORS does not apply to websockets, so check the origin ourselves.
|
|
if websocket.headers.get("origin") not in origins:
|
|
await websocket.close(code=1008)
|
|
return
|
|
if (sum(active_websockets.values()) >= MAX_WEBSOCKETS
|
|
or active_websockets[ip] >= MAX_WEBSOCKETS_PER_IP):
|
|
await websocket.close(code=1013) # try again later
|
|
return
|
|
|
|
active_websockets[ip] += 1
|
|
try:
|
|
await websocket.accept()
|
|
last = 0.0
|
|
|
|
while True:
|
|
message = await websocket.receive()
|
|
if message["type"] == "websocket.disconnect":
|
|
break
|
|
|
|
data = message.get("bytes")
|
|
if data is None or len(data) > MAX_FRAME_BYTES:
|
|
await websocket.close(code=1009 if data else 1003)
|
|
break
|
|
|
|
# Drop frames that arrive too fast or while the server is busy
|
|
# instead of queueing them.
|
|
now = time.monotonic()
|
|
if now - last < FACE_MIN_INTERVAL:
|
|
continue
|
|
last = now
|
|
|
|
try:
|
|
boxes = await gate.run(detect_faces, data)
|
|
except Busy:
|
|
continue
|
|
except InvalidImage:
|
|
await websocket.close(code=1003)
|
|
break
|
|
|
|
await websocket.send_json({"bboxes": boxes})
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception:
|
|
logger.exception("predict_face failed")
|
|
try:
|
|
await websocket.close(code=1011)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
active_websockets[ip] -= 1
|
|
if active_websockets[ip] <= 0:
|
|
del active_websockets[ip]
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Reviews
|
|
# --------------------------------------------------------------------------
|
|
|
|
def get_api_key(authorization: str = Header(None)):
|
|
if not API_KEY or not authorization or not secrets.compare_digest(
|
|
authorization.encode(), f"Bearer {API_KEY}".encode()
|
|
):
|
|
raise HTTPException(status_code=401, detail="Unauthorized")
|
|
return authorization
|
|
|
|
|
|
class Review(BaseModel):
|
|
website: str = Field(min_length=1, max_length=100)
|
|
rating: int = Field(ge=1, le=5)
|
|
text: str = Field(max_length=1000)
|
|
# The client also sends a "date"; it is ignored, the server sets the time.
|
|
|
|
|
|
def connect_db() -> sqlite3.Connection:
|
|
conn = sqlite3.connect(DATABASE, timeout=5)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
|
|
@app.get("/review")
|
|
def get_reviews(
|
|
limit: int = Query(100, ge=1, le=500),
|
|
offset: int = Query(0, ge=0),
|
|
auth=Depends(get_api_key),
|
|
):
|
|
with closing(connect_db()) as conn:
|
|
rows = conn.execute(
|
|
"SELECT * FROM reviews ORDER BY id DESC LIMIT ? OFFSET ?", (limit, offset)
|
|
).fetchall()
|
|
return {"data": [dict(row) for row in rows]}
|
|
|
|
|
|
# Public on purpose: no credentials needed, only rate limited per IP.
|
|
@app.post("/review", dependencies=[Depends(rate_limit(review_limiter))])
|
|
def post_review(review: Review):
|
|
created_at = datetime.now(timezone.utc).isoformat()
|
|
with closing(connect_db()) as conn:
|
|
conn.execute(
|
|
"INSERT INTO reviews (website, rating, text, created_at) VALUES (?, ?, ?, ?)",
|
|
(review.website, review.rating, review.text.strip(), created_at),
|
|
)
|
|
conn.commit()
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.delete("/review")
|
|
def delete_review(auth=Depends(get_api_key)):
|
|
with closing(connect_db()) as conn:
|
|
conn.execute("DELETE FROM reviews")
|
|
conn.commit()
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|