Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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 [], | |
| } | |
| 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) | |
| } | |
| def status(): | |
| return training_status() | |
| 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'} | |
| 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)) | |
| 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.') | |