Upload 3 files

README.md

@@ -0,0 +1,19 @@
+# Orcan VisionTrace GPU Service
+
+GPU-accelerated face recognition and FAISS indexing service for Orcan VisionTrace.
+
+## Features
+- Batch face embedding extraction using InsightFace
+- GPU-accelerated FAISS index creation
+- Image enhancement for poor quality inputs
+- High-performance search capabilities
+
+## Hardware Requirements
+- NVIDIA GPU with CUDA support
+- Minimum 8GB GPU memory recommended
+
+## API Endpoints
+- POST /extract_embeddings_batch - Batch face embedding extraction
+- POST /create_faiss_index - GPU-accelerated index creation  
+- POST /search_faiss - Fast similarity search
+- GET /health - Service health check

app.py

@@ -0,0 +1,222 @@
+import os
+import base64
+import numpy as np
+import cv2
+import faiss
+import torch
+import insightface
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Dict, Any, Optional
+from PIL import Image, ImageOps
+import io
+import logging
+
+app = FastAPI(title="Orcan VisionTrace GPU Service", version="1.0.0")
+
+# Global models
+face_app = None
+gpu_resources = None
+
+class BatchEmbeddingRequest(BaseModel):
+    images: List[str]  # Base64 encoded images
+    enhance_quality: bool = True
+    aggressive_enhancement: bool = False
+
+class IndexCreationRequest(BaseModel):
+    embeddings: List[List[float]]
+    dataset_size: int
+    dimension: int = 512
+
+@app.on_event("startup")
+async def startup_event():
+    global face_app, gpu_resources
+    
+    # Initialize FAISS GPU resources
+    if torch.cuda.is_available():
+        gpu_resources = faiss.StandardGpuResources()
+        
+    # Initialize InsightFace with GPU
+    face_app = insightface.app.FaceAnalysis(
+        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
+        allowed_modules=['detection', 'recognition']
+    )
+    face_app.prepare(ctx_id=0, det_size=(640, 640))
+
+@app.get("/health")
+async def health_check():
+    return {
+        "status": "healthy",
+        "gpu_available": torch.cuda.is_available(),
+        "face_model_loaded": face_app is not None
+    }
+
+@app.post("/extract_embeddings_batch")
+async def extract_embeddings_batch(request: BatchEmbeddingRequest):
+    try:
+        embeddings = []
+        extraction_info = []
+        
+        for img_b64 in request.images:
+            try:
+                # Decode base64 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 image"})
+                    continue
+                
+                # Apply enhancement if requested
+                if request.enhance_quality:
+                    img = enhance_image_gpu(img, request.aggressive_enhancement)
+                
+                # Extract face embeddings
+                faces = face_app.get(img)
+                
+                if len(faces) == 0:
+                    embeddings.append(None)
+                    extraction_info.append({
+                        "face_count": 0,
+                        "strategy_used": "gpu_batch",
+                        "enhancement_used": request.enhance_quality
+                    })
+                    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
+                embedding = embedding / np.linalg.norm(embedding)
+                
+                embeddings.append(embedding.tolist())
+                extraction_info.append({
+                    "face_count": len(faces),
+                    "confidence": float(np.linalg.norm(embedding)),
+                    "strategy_used": "gpu_batch",
+                    "enhancement_used": request.enhance_quality,
+                    "quality_score": 0.8  # Placeholder
+                })
+                
+            except Exception as e:
+                embeddings.append(None)
+                extraction_info.append({"error": str(e)})
+        
+        return {
+            "embeddings": embeddings,
+            "extraction_info": extraction_info,
+            "total_processed": len(request.images),
+            "successful": len([e for e in embeddings if e is not None])
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+def enhance_image_gpu(img, aggressive=False):
+    """GPU-accelerated image enhancement"""
+    if aggressive:
+        # Strong enhancement for poor quality images
+        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)
+        
+        # Sharpening
+        kernel = np.array([[-1,-1,-1], [-1, 12,-1], [-1,-1,-1]])
+        img = cv2.filter2D(img, -1, kernel)
+    else:
+        # Standard enhancement
+        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
+
+@app.post("/create_faiss_index")
+async def create_faiss_index(request: IndexCreationRequest):
+    try:
+        embeddings_array = np.array(request.embeddings, dtype='float32')
+        
+        # Choose index type based on dataset size
+        if request.dataset_size < 1000:
+            index = faiss.IndexFlatL2(request.dimension)
+            index_type = "IndexFlatL2"
+        elif request.dataset_size < 50000:
+            nlist = max(4, min(request.dataset_size // 39, 100))
+            quantizer = faiss.IndexFlatL2(request.dimension)
+            index = faiss.IndexIVFFlat(quantizer, request.dimension, nlist)
+            index_type = "IndexIVFFlat"
+        else:
+            nlist = max(100, min(request.dataset_size // 39, 1000))
+            quantizer = faiss.IndexFlatL2(request.dimension)
+            index = faiss.IndexIVFPQ(quantizer, request.dimension, nlist, 64, 8)
+            index_type = "IndexIVFPQ"
+        
+        # Move to GPU if available
+        if torch.cuda.is_available() and gpu_resources is not None:
+            if hasattr(index, 'train'):
+                # Train on GPU
+                index_gpu = faiss.index_cpu_to_gpu(gpu_resources, 0, index)
+                if not index_gpu.is_trained:
+                    index_gpu.train(embeddings_array)
+                index_gpu.add(embeddings_array)
+                # Move back to CPU for serialization
+                index = faiss.index_gpu_to_cpu(index_gpu)
+            else:
+                # Flat index - direct GPU processing
+                index_gpu = faiss.index_cpu_to_gpu(gpu_resources, 0, index)
+                index_gpu.add(embeddings_array)
+                index = faiss.index_gpu_to_cpu(index_gpu)
+        else:
+            # CPU fallback
+            if hasattr(index, 'train') and not index.is_trained:
+                index.train(embeddings_array)
+            index.add(embeddings_array)
+        
+        # Serialize index
+        index_data = faiss.serialize_index(index)
+        index_b64 = base64.b64encode(index_data).decode()
+        
+        return {
+            "index_data": index_b64,
+            "index_type": f"GPU_{index_type}",
+            "index_params": {"nlist": getattr(index, 'nlist', 0)},
+            "vectors_added": index.ntotal
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/search_faiss")
+async def search_faiss(request: dict):
+    try:
+        # Deserialize index
+        index_data = base64.b64decode(request["index_data"])
+        index = faiss.deserialize_index(np.frombuffer(index_data, dtype=np.uint8))
+        
+        query_embedding = np.array([request["query_embedding"]], dtype='float32')
+        k = request.get("k", 25)
+        
+        # Move to GPU for search
+        if torch.cuda.is_available() and gpu_resources is not None:
+            index_gpu = faiss.index_cpu_to_gpu(gpu_resources, 0, index)
+            distances, indices = index_gpu.search(query_embedding, k)
+        else:
+            distances, indices = index.search(query_embedding, k)
+        
+        return {
+            "distances": distances[0].tolist(),
+            "indices": indices[0].tolist(),
+            "total_vectors": index.ntotal
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -0,0 +1,12 @@
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+torch==2.1.0
+torchvision==0.16.0
+faiss-gpu==1.7.4
+insightface==0.7.3
+opencv-python==4.8.1.78
+Pillow==10.1.0
+numpy==1.24.3
+pydantic==2.4.2
+python-multipart==0.0.6
+onnxruntime-gpu==1.16.0