Spaces:
Runtime error
Runtime error
File size: 6,188 Bytes
48922fa |
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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 |
"""
Image analysis engine for processing and analyzing images.
"""
from typing import Dict, Any, List, Optional
import io
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import face_recognition
import numpy as np
from tenacity import retry, stop_after_attempt, wait_exponential
class ImageEngine:
def __init__(self):
# Initialize image classification model
self.feature_extractor = AutoFeatureExtractor.from_pretrained(
"microsoft/resnet-50"
)
self.model = AutoModelForImageClassification.from_pretrained(
"microsoft/resnet-50"
)
# Set up image transforms
self.transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
async def analyze_image(self, image_data: bytes) -> Dict[str, Any]:
"""Analyze image content and detect objects/faces."""
try:
# Load image
image = Image.open(io.BytesIO(image_data)).convert('RGB')
# Prepare image for model
inputs = self.feature_extractor(images=image, return_tensors="pt")
# Get model predictions
with torch.no_grad():
outputs = self.model(**inputs)
probs = outputs.logits.softmax(-1)
# Get top predictions
top_probs, top_indices = torch.topk(probs, k=5)
# Convert predictions to list
predictions = [
{
"label": self.model.config.id2label[idx.item()],
"confidence": prob.item()
}
for prob, idx in zip(top_probs[0], top_indices[0])
]
# Analyze faces
np_image = np.array(image)
face_locations = face_recognition.face_locations(np_image)
face_encodings = face_recognition.face_encodings(np_image, face_locations)
faces = []
for i, (face_encoding, face_location) in enumerate(zip(face_encodings, face_locations)):
face = {
"id": i + 1,
"location": {
"top": face_location[0],
"right": face_location[1],
"bottom": face_location[2],
"left": face_location[3]
},
"encoding": face_encoding.tolist()
}
faces.append(face)
# Get image metadata
metadata = {
"format": image.format,
"mode": image.mode,
"size": image.size,
"width": image.width,
"height": image.height
}
return {
"predictions": predictions,
"faces": faces,
"metadata": metadata
}
except Exception as e:
return {"error": str(e)}
async def compare_faces(self, face1_data: bytes, face2_data: bytes) -> Dict[str, Any]:
"""Compare two faces and determine if they are the same person."""
try:
# Load and process first image
image1 = face_recognition.load_image_file(io.BytesIO(face1_data))
face1_encoding = face_recognition.face_encodings(image1)
if not face1_encoding:
return {"error": "No face found in first image"}
# Load and process second image
image2 = face_recognition.load_image_file(io.BytesIO(face2_data))
face2_encoding = face_recognition.face_encodings(image2)
if not face2_encoding:
return {"error": "No face found in second image"}
# Compare faces
results = face_recognition.compare_faces(
[face1_encoding[0]], face2_encoding[0]
)
# Calculate face distance (lower means more similar)
face_distance = face_recognition.face_distance(
[face1_encoding[0]], face2_encoding[0]
)
return {
"match": bool(results[0]),
"confidence": float(1 - face_distance[0]),
"distance": float(face_distance[0])
}
except Exception as e:
return {"error": str(e)}
async def search_similar_faces(self,
target_encoding: List[float],
face_database: List[Dict[str, Any]],
threshold: float = 0.6) -> List[Dict[str, Any]]:
"""Search for similar faces in a database of face encodings."""
try:
matches = []
target_encoding = np.array(target_encoding)
for face_data in face_database:
if "encoding" not in face_data:
continue
current_encoding = np.array(face_data["encoding"])
distance = face_recognition.face_distance([target_encoding], current_encoding)[0]
if distance < threshold:
matches.append({
"face_id": face_data.get("id"),
"confidence": float(1 - distance),
"metadata": face_data.get("metadata", {})
})
# Sort matches by confidence
matches.sort(key=lambda x: x["confidence"], reverse=True)
return matches
except Exception as e:
return [{"error": str(e)}]
|