Spaces:
Running on Zero
Running on Zero
File size: 3,286 Bytes
f2798b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import json
import os
import shutil
import threading
import uuid
from pathlib import Path
from typing import Optional
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
from paddle_backend import train, training_status, paddle_status, predict, model_ready
ROOT = Path(__file__).resolve().parent
if os.getenv('DATA_DIR'):
BASE = Path(os.environ['DATA_DIR'])
elif Path('/data').exists() and os.access('/data', os.W_OK):
BASE = Path('/data') / 'icecream_counter'
else:
BASE = Path('./data')
DATASET = BASE / 'generated_dataset'
MODEL = BASE / 'model'
CLASSES = BASE / 'classes.txt'
API_OUTPUT = BASE / 'api_outputs'
API_OUTPUT.mkdir(parents=True, exist_ok=True)
app = FastAPI(title='Ice Cream RT-DETR API', version='1.0.0')
class TrainRequest(BaseModel):
epochs: int = 30
batch_size: int = 2
learning_rate: float = 1e-5
@app.get('/dataset/status')
def dataset_status():
return {
'train_annotations': (DATASET / 'train' / 'annotations.json').exists(),
'val_annotations': (DATASET / 'val' / 'annotations.json').exists(),
'classes': [x.strip() for x in CLASSES.read_text(encoding='utf-8').splitlines() if x.strip()] if CLASSES.exists() else [],
}
@app.get('/health')
def health():
p = paddle_status()
return {
'ok': True,
'engine': 'PaddleDetection RT-DETR',
'huggingface': False,
'transformers': False,
'torch_training': False,
'paddle': p,
'model_ready': model_ready(MODEL)
}
@app.get('/train/status')
def status():
return training_status()
@app.post('/train')
def start_train(req: TrainRequest):
s = training_status()
if s['running']:
return s
if not (DATASET / 'train' / 'annotations.json').exists():
raise HTTPException(400, 'Training dataset is not prepared.')
if not CLASSES.exists():
raise HTTPException(400, 'classes.txt is missing.')
def job():
train(DATASET, CLASSES, MODEL, max(1, req.epochs), max(1, req.batch_size), req.learning_rate)
threading.Thread(target=job, daemon=True).start()
return {'running': True, 'message': 'training started', 'engine': 'PaddleDetection RT-DETR'}
@app.post('/predict')
async def predict_image(file: UploadFile = File(...), threshold: float = 0.35):
if not model_ready(MODEL):
raise HTTPException(409, 'No trained model yet.')
job_id = uuid.uuid4().hex
job_dir = API_OUTPUT / job_id
job_dir.mkdir(parents=True, exist_ok=True)
suffix = Path(file.filename or 'image.jpg').suffix or '.jpg'
src = job_dir / ('input' + suffix)
src.write_bytes(await file.read())
try:
result = predict(src, MODEL, max(0.0, min(1.0, threshold)), CLASSES, job_dir)
result['job_id'] = job_id
return result
except Exception as e:
shutil.rmtree(job_dir, ignore_errors=True)
raise HTTPException(500, str(e))
@app.get('/predict/{job_id}/image')
def predicted_image(job_id: str):
d = API_OUTPUT / job_id
for p in d.glob('*'):
if p.suffix.lower() in {'.jpg','.jpeg','.png'} and p.name != 'input.jpg':
return FileResponse(p)
raise HTTPException(404, 'Prediction image not found.')
|