File size: 5,551 Bytes
ff08aee 9a7a94e 0fce058 ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e ff08aee 9a7a94e |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
from typing import Dict, List, Any
import json
import base64
import numpy as np
import cv2
import torch
import insightface
from PIL import Image
import io
class EndpointHandler:
def __init__(self, path=""):
self.face_app = None
self.use_gpu = False
self._init_model()
def _init_model(self):
"""Initialize InsightFace model"""
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
ctx_id = 0
else:
providers = ['CPUExecutionProvider']
ctx_id = -1
self.face_app = insightface.app.FaceAnalysis(
providers=providers,
allowed_modules=['detection', 'recognition']
)
self.face_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
print(f"Face model loaded: {'GPU' if self.use_gpu else 'CPU'}")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""Handle the actual inference request"""
try:
# Handle health check
if data.get("inputs") == "test":
return {
"status": "healthy",
"gpu_available": self.use_gpu,
"model_loaded": self.face_app is not None
}
# Handle batch embedding extraction
if "images" in data:
return self._extract_embeddings_batch(data)
return {"error": "Unknown request format"}
except Exception as e:
return {"error": str(e)}
def _extract_embeddings_batch(self, data):
"""Extract embeddings from batch of images"""
images = data.get("images", [])
enhance_quality = data.get("enhance_quality", True)
aggressive = data.get("aggressive_enhancement", False)
embeddings = []
extraction_info = []
for idx, img_b64 in enumerate(images):
try:
# Decode image
img_data = base64.b64decode(img_b64)
img_array = np.frombuffer(img_data, dtype=np.uint8)
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
if img is None:
embeddings.append(None)
extraction_info.append({"error": "Failed to decode", "index": idx})
continue
# Enhance if requested
if enhance_quality:
img = self._enhance_image(img, aggressive)
# Extract faces
faces = self.face_app.get(img)
if len(faces) == 0:
embeddings.append(None)
extraction_info.append({
"face_count": 0,
"strategy_used": "gpu_batch" if self.use_gpu else "cpu_batch",
"enhancement_used": enhance_quality,
"index": idx
})
continue
# Get best face
face = max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
embedding = face.embedding / np.linalg.norm(face.embedding)
embeddings.append(embedding.tolist())
# Calculate metrics
bbox_area = (face.bbox[2] - face.bbox[0]) * (face.bbox[3] - face.bbox[1])
img_area = img.shape[0] * img.shape[1]
confidence = min((bbox_area / img_area) * 2.0, 1.0)
extraction_info.append({
"face_count": len(faces),
"confidence": float(confidence),
"strategy_used": "gpu_batch" if self.use_gpu else "cpu_batch",
"enhancement_used": enhance_quality,
"quality_score": float(confidence),
"index": idx
})
except Exception as e:
embeddings.append(None)
extraction_info.append({"error": str(e), "index": idx})
successful = len([e for e in embeddings if e is not None])
return {
"embeddings": embeddings,
"extraction_info": extraction_info,
"total_processed": len(images),
"successful": successful,
"processing_mode": "gpu" if self.use_gpu else "cpu"
}
def _enhance_image(self, img, aggressive=False):
"""Image enhancement logic"""
try:
if aggressive:
img = cv2.bilateralFilter(img, 15, 90, 90)
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
l, a, b = cv2.split(lab)
clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(8,8))
l = clahe.apply(l)
img = cv2.merge([l, a, b])
img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
else:
img = cv2.bilateralFilter(img, 9, 75, 75)
kernel = np.array([[-1,-1,-1], [-1, 9,-1], [-1,-1,-1]])
img = cv2.filter2D(img, -1, kernel)
return img
except:
return img |