You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
175 lines
5.6 KiB
175 lines
5.6 KiB
from contextlib import asynccontextmanager
|
|
import io
|
|
from typing import List
|
|
import numpy as np
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from PIL import Image
|
|
from pydantic import BaseModel
|
|
import insightface
|
|
from insightface.app import FaceAnalysis
|
|
from face_store import FaceStore
|
|
import time
|
|
|
|
# Initializing InsightFace model with better detection settings with robust detection
|
|
face_analyzer = FaceAnalysis(
|
|
providers=['CPUExecutionProvider'],
|
|
allowed_modules=['detection', 'recognition']
|
|
)
|
|
face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
|
|
|
|
# Initializing face store
|
|
face_store = FaceStore()
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
print("Initializing face recognition system...")
|
|
yield
|
|
print("Cleaning up resources...")
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
class Visitor(BaseModel):
|
|
name: str
|
|
encoded_face: List[float]
|
|
|
|
class Config:
|
|
from_attributes = True
|
|
|
|
def get_largest_face(faces):
|
|
"""Select the largest face from detected faces based on bounding box area."""
|
|
if not faces:
|
|
return None
|
|
|
|
# Calculating areas of all faces
|
|
areas = [(face, (face.bbox[2] - face.bbox[0]) * (face.bbox[3] - face.bbox[1]))
|
|
for face in faces]
|
|
# Returning face with largest area
|
|
return max(areas, key=lambda x: x[1])[0]
|
|
|
|
def process_image(image_data: bytes):
|
|
"""Process image and return embedding of the largest face."""
|
|
try:
|
|
image_stream = io.BytesIO(image_data)
|
|
image_pil = Image.open(image_stream).convert("RGB")
|
|
|
|
# Resizing image if too large (optional, adjust dimensions as needed)
|
|
max_size = 1920
|
|
if max(image_pil.size) > max_size:
|
|
ratio = max_size / max(image_pil.size)
|
|
new_size = tuple(int(dim * ratio) for dim in image_pil.size)
|
|
image_pil = image_pil.resize(new_size, Image.Resampling.LANCZOS)
|
|
|
|
image_np = np.array(image_pil)
|
|
faces = face_analyzer.get(image_np)
|
|
|
|
if not faces:
|
|
return None, "No face detected"
|
|
|
|
# Get the largest face
|
|
largest_face = get_largest_face(faces)
|
|
|
|
# Converting embedding to numpy array to ensure consistent format
|
|
embedding = np.array(largest_face.embedding, dtype=np.float32)
|
|
return embedding, None
|
|
|
|
except Exception as e:
|
|
return None, f"Error processing image: {str(e)}"
|
|
|
|
@app.get("/")
|
|
async def health_check():
|
|
return {"message": "Face recognition API is running"}
|
|
|
|
@app.post("/api/register")
|
|
async def register_visitor(name: str = Form(...), image: UploadFile = File(...)):
|
|
try:
|
|
image_data = await image.read()
|
|
embedding, error = process_image(image_data)
|
|
|
|
if error:
|
|
return {"message": error}
|
|
|
|
# Converting embedding to numpy array if it isn't already
|
|
embedding = np.array(embedding, dtype=np.float32)
|
|
|
|
# Adding debug logging
|
|
print(f"Registering face for {name}")
|
|
print(f"Embedding shape: {embedding.shape}")
|
|
print(f"Embedding type: {type(embedding)}")
|
|
|
|
# Checking if face already exists
|
|
existing_match = face_store.search_face(embedding)
|
|
if existing_match:
|
|
return {
|
|
"message": "Visitor already exists",
|
|
"name": existing_match[0]
|
|
}
|
|
|
|
# Registering new face
|
|
face_store.add_face(name, embedding)
|
|
|
|
# Verifying registration
|
|
verification = face_store.search_face(embedding)
|
|
if not verification:
|
|
raise HTTPException(status_code=500, detail="Face registration failed verification")
|
|
|
|
return {
|
|
"message": "Visitor registered successfully",
|
|
"name": name
|
|
}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app.post("/api/search")
|
|
async def search_visitor(image: UploadFile = File(...)):
|
|
try:
|
|
image_data = await image.read()
|
|
start_time = time.time()
|
|
print(f"API request started at: {time.strftime('%H:%M:%S')}")
|
|
|
|
embedding, error = process_image(image_data)
|
|
|
|
if error:
|
|
return {"message": error}
|
|
|
|
# Converting embedding to numpy array if it isn't already
|
|
embedding = np.array(embedding, dtype=np.float32)
|
|
|
|
# Adding debug logging
|
|
print(f"Searching for face")
|
|
print(f"Embedding shape: {embedding.shape}")
|
|
print(f"Embedding type: {type(embedding)}")
|
|
|
|
search_start = time.time()
|
|
match = face_store.search_face(embedding)
|
|
search_time = time.time() - search_start
|
|
print(f"Face search took: {search_time:.4f} seconds")
|
|
|
|
total_time = time.time() - start_time
|
|
print(f"Total API processing time: {total_time:.4f} seconds")
|
|
|
|
if match:
|
|
name, confidence = match
|
|
return {
|
|
"message": "Visitor found",
|
|
"name": name,
|
|
"confidence": confidence,
|
|
"search_time": search_time,
|
|
"processing_time": total_time
|
|
}
|
|
return {
|
|
"message": "Visitor not found",
|
|
"processing_time": total_time
|
|
}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|