133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
from fastapi import FastAPI, File, UploadFile
|
|
from fastapi.responses import JSONResponse
|
|
from anime_face_detector import create_detector
|
|
import numpy as np
|
|
import cv2
|
|
import os
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
app = FastAPI(title="Anime Face Detector API")
|
|
|
|
# Create output directory
|
|
OUTPUT_DIR = Path("/app/api/outputs")
|
|
OUTPUT_DIR.mkdir(exist_ok=True)
|
|
|
|
# Load detector once at startup
|
|
detector = None
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
global detector
|
|
detector = create_detector("yolov3", device="cuda:0")
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
return {"message": "Anime Face Detector API", "status": "running"}
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "healthy", "detector_loaded": detector is not None}
|
|
|
|
@app.post("/detect")
|
|
async def detect_face(file: UploadFile = File(...)):
|
|
try:
|
|
if detector is None:
|
|
return JSONResponse(
|
|
status_code=503,
|
|
content={"error": "Detector not initialized yet"}
|
|
)
|
|
|
|
image_bytes = await file.read()
|
|
|
|
# Decode image
|
|
nparr = np.frombuffer(image_bytes, np.uint8)
|
|
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
|
|
if img is None:
|
|
return JSONResponse(
|
|
status_code=400,
|
|
content={"error": "Invalid image file"}
|
|
)
|
|
|
|
# Run detector
|
|
preds = detector(img)
|
|
|
|
# Convert output to JSON-safe structure
|
|
result = []
|
|
for pred in preds:
|
|
# bbox contains [x1, y1, x2, y2, score]
|
|
bbox = pred['bbox']
|
|
bbox_coords = bbox[:4].tolist()
|
|
score = float(bbox[4])
|
|
|
|
# keypoints contains [[x, y, score], ...]
|
|
keypoints = pred.get('keypoints', [])
|
|
landmarks = [[float(pt[0]), float(pt[1]), float(pt[2])] for pt in keypoints]
|
|
|
|
result.append({
|
|
"bbox": bbox_coords,
|
|
"confidence": score,
|
|
"keypoints": landmarks
|
|
})
|
|
|
|
# Draw detection results on image (like Gradio demo)
|
|
annotated_img = img.copy()
|
|
for pred in preds:
|
|
box = pred['bbox']
|
|
box_coords, score = box[:4], box[4]
|
|
box_coords = np.round(box_coords).astype(int)
|
|
|
|
# Calculate line thickness based on image size
|
|
lt = max(2, int(3 * (box_coords[2:] - box_coords[:2]).max() / 256))
|
|
|
|
# Draw bounding box
|
|
cv2.rectangle(annotated_img,
|
|
tuple(box_coords[:2]),
|
|
tuple(box_coords[2:]),
|
|
(0, 255, 0), lt)
|
|
|
|
# Draw keypoints
|
|
pred_pts = pred['keypoints']
|
|
for *pt, kp_score in pred_pts:
|
|
# Color based on keypoint confidence
|
|
if kp_score < 0.3: # Low confidence
|
|
color = (0, 255, 255) # Yellow
|
|
else:
|
|
color = (0, 0, 255) # Red
|
|
pt = np.round(pt).astype(int)
|
|
cv2.circle(annotated_img, tuple(pt), lt, color, cv2.FILLED)
|
|
|
|
# Generate unique filename based on timestamp
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
|
original_filename = Path(file.filename).stem if file.filename else "image"
|
|
|
|
# Save annotated image
|
|
image_output_path = OUTPUT_DIR / f"{original_filename}_{timestamp}_annotated.jpg"
|
|
cv2.imwrite(str(image_output_path), annotated_img)
|
|
|
|
# Save JSON results
|
|
json_output_path = OUTPUT_DIR / f"{original_filename}_{timestamp}_results.json"
|
|
json_data = {
|
|
"filename": file.filename,
|
|
"timestamp": timestamp,
|
|
"detections": result,
|
|
"count": len(result)
|
|
}
|
|
with open(json_output_path, 'w') as f:
|
|
json.dump(json_data, f, indent=2)
|
|
|
|
return {
|
|
"detections": result,
|
|
"count": len(result),
|
|
"annotated_image": str(image_output_path),
|
|
"json_file": str(json_output_path)
|
|
}
|
|
|
|
except Exception as e:
|
|
return JSONResponse(
|
|
status_code=500,
|
|
content={"error": str(e)}
|
|
)
|