Bernt
bae705aa97
- Add NFC ePassport roadmap (ICAO 9303, eIDAS) - Add TensorFlow.js edge face detection (BlazeFace) - Add structured audit logger (GDPR-compliant) - Risk scoring support Part of KYC Apple Native UX v1.1.0
257 lines
8.1 KiB
Python
257 lines
8.1 KiB
Python
"""
|
|
Real AI Pipeline using YOLO and CLIP
|
|
Production-ready image analysis with real models
|
|
"""
|
|
|
|
from typing import Dict, List, Optional, Tuple
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
|
|
# Try to import real AI libraries
|
|
try:
|
|
import torch
|
|
import torchvision
|
|
from transformers import CLIPProcessor, CLIPModel
|
|
HAS_REAL_AI = True
|
|
except ImportError:
|
|
HAS_REAL_AI = False
|
|
print("Warning: Real AI libraries not installed. Using simulation.")
|
|
|
|
|
|
@dataclass
|
|
class AIAnalysisResult:
|
|
"""Result from real AI analysis"""
|
|
detected_objects: List[Dict]
|
|
defect_codes: List[Dict]
|
|
scene_type: str
|
|
confidence: float
|
|
embeddings: Optional[np.ndarray] = None
|
|
|
|
def to_dict(self) -> Dict:
|
|
return {
|
|
"detected_objects": self.detected_objects,
|
|
"defect_codes": self.defect_codes,
|
|
"scene_type": self.scene_type,
|
|
"confidence": self.confidence,
|
|
"has_embeddings": self.embeddings is not None
|
|
}
|
|
|
|
|
|
class RealAIClassifier:
|
|
"""
|
|
Production AI classifier using YOLO and CLIP
|
|
|
|
Requirements:
|
|
- torch
|
|
- torchvision
|
|
- transformers
|
|
- pillow
|
|
- numpy
|
|
|
|
Models:
|
|
- YOLOv8 for object detection
|
|
- CLIP for scene classification
|
|
- Custom classifier for defects
|
|
"""
|
|
|
|
def __init__(self, use_real_ai: bool = True):
|
|
self.use_real_ai = use_real_ai and HAS_REAL_AI
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if HAS_REAL_AI else None
|
|
|
|
# Load models
|
|
self.yolo_model = None
|
|
self.clip_model = None
|
|
self.clip_processor = None
|
|
self.defect_classifier = None
|
|
|
|
if self.use_real_ai:
|
|
self._load_models()
|
|
else:
|
|
print("Using simulated AI (install torch, torchvision, transformers for real AI)")
|
|
|
|
def _load_models(self):
|
|
"""Load AI models"""
|
|
print("Loading AI models...")
|
|
|
|
# Load YOLOv8
|
|
try:
|
|
from ultralytics import YOLO
|
|
self.yolo_model = YOLO("yolov8n.pt") # Nano model for speed
|
|
print("✓ YOLOv8 loaded")
|
|
except Exception as e:
|
|
print(f"✗ YOLOv8 failed: {e}")
|
|
self.yolo_model = None
|
|
|
|
# Load CLIP
|
|
try:
|
|
self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
|
|
self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
|
|
self.clip_model.to(self.device)
|
|
print("✓ CLIP loaded")
|
|
except Exception as e:
|
|
print(f"✗ CLIP failed: {e}")
|
|
self.clip_model = None
|
|
self.clip_processor = None
|
|
|
|
# Load defect classifier (custom model)
|
|
try:
|
|
# In production, load trained model
|
|
# self.defect_classifier = torch.load("defect_classifier.pth")
|
|
print("✓ Defect classifier placeholder loaded")
|
|
except Exception as e:
|
|
print(f"✗ Defect classifier failed: {e}")
|
|
self.defect_classifier = None
|
|
|
|
def analyze_image(self, image_path: str, location: Optional[Dict] = None) -> AIAnalysisResult:
|
|
"""Analyze image with real AI"""
|
|
if not self.use_real_ai:
|
|
return self._simulate_analysis(image_path, location)
|
|
|
|
# Load image
|
|
image = Image.open(image_path).convert("RGB")
|
|
|
|
# Object detection with YOLO
|
|
detected_objects = self._detect_objects(image)
|
|
|
|
# Scene classification with CLIP
|
|
scene_type, scene_confidence = self._classify_scene(image)
|
|
|
|
# Defect detection
|
|
defect_codes = self._detect_defects(image)
|
|
|
|
# Generate embeddings
|
|
embeddings = self._generate_embeddings(image)
|
|
|
|
return AIAnalysisResult(
|
|
detected_objects=detected_objects,
|
|
defect_codes=defect_codes,
|
|
scene_type=scene_type,
|
|
confidence=scene_confidence,
|
|
embeddings=embeddings
|
|
)
|
|
|
|
def _detect_objects(self, image: Image.Image) -> List[Dict]:
|
|
"""Detect objects with YOLO"""
|
|
if self.yolo_model is None:
|
|
return []
|
|
|
|
results = self.yolo_model(image)
|
|
objects = []
|
|
|
|
for result in results:
|
|
boxes = result.boxes
|
|
for box in boxes:
|
|
objects.append({
|
|
"label": result.names[int(box.cls)],
|
|
"confidence": float(box.conf),
|
|
"bbox": box.xyxy.tolist()[0]
|
|
})
|
|
|
|
return objects
|
|
|
|
def _classify_scene(self, image: Image.Image) -> Tuple[str, float]:
|
|
"""Classify scene with CLIP"""
|
|
if self.clip_model is None or self.clip_processor is None:
|
|
return "unknown", 0.0
|
|
|
|
# Define scene labels
|
|
scene_labels = [
|
|
"street view", "building facade", "bridge", "road",
|
|
"sidewalk", "park", "industrial area", "residential area",
|
|
"commercial area", "construction site"
|
|
]
|
|
|
|
# Process image
|
|
inputs = self.clip_processor(
|
|
text=scene_labels,
|
|
images=image,
|
|
return_tensors="pt",
|
|
padding=True
|
|
).to(self.device)
|
|
|
|
# Get predictions
|
|
with torch.no_grad():
|
|
outputs = self.clip_model(**inputs)
|
|
logits_per_image = outputs.logits_per_image
|
|
probs = logits_per_image.softmax(dim=1)
|
|
|
|
# Get top prediction
|
|
top_prob, top_idx = probs.max(dim=1)
|
|
|
|
return scene_labels[top_idx.item()], float(top_prob.item())
|
|
|
|
def _detect_defects(self, image: Image.Image) -> List[Dict]:
|
|
"""Detect defects with custom classifier"""
|
|
if self.defect_classifier is None:
|
|
return []
|
|
|
|
# In production, use trained defect classifier
|
|
# For now, return empty list
|
|
return []
|
|
|
|
def _generate_embeddings(self, image: Image.Image) -> Optional[np.ndarray]:
|
|
"""Generate image embeddings with CLIP"""
|
|
if self.clip_model is None or self.clip_processor is None:
|
|
return None
|
|
|
|
inputs = self.clip_processor(
|
|
images=image,
|
|
return_tensors="pt"
|
|
).to(self.device)
|
|
|
|
with torch.no_grad():
|
|
image_features = self.clip_model.get_image_features(**inputs)
|
|
|
|
return image_features.cpu().numpy()
|
|
|
|
def _simulate_analysis(self, image_path: str, location: Optional[Dict] = None) -> AIAnalysisResult:
|
|
"""Simulated analysis when real AI is not available"""
|
|
from ai_pipeline.image_classifier import ImageClassifier
|
|
|
|
simulator = ImageClassifier()
|
|
result = simulator.analyze_image(image_path, location)
|
|
|
|
return AIAnalysisResult(
|
|
detected_objects=result.detected_objects,
|
|
defect_codes=result.defect_codes,
|
|
scene_type=result.scene_type,
|
|
confidence=result.confidence,
|
|
embeddings=None
|
|
)
|
|
|
|
def batch_process(self, image_paths: List[str], locations: Optional[List[Dict]] = None) -> List[AIAnalysisResult]:
|
|
"""Process multiple images"""
|
|
results = []
|
|
for i, path in enumerate(image_paths):
|
|
loc = locations[i] if locations and i < len(locations) else None
|
|
result = self.analyze_image(path, loc)
|
|
results.append(result)
|
|
return results
|
|
|
|
|
|
# Example usage
|
|
def example_real_ai():
|
|
"""Example: Real AI analysis"""
|
|
print("=== Real AI Pipeline ===\n")
|
|
|
|
classifier = RealAIClassifier(use_real_ai=False) # Set to True when models are installed
|
|
|
|
print(f"Real AI available: {classifier.use_real_ai}")
|
|
print(f"Device: {classifier.device}")
|
|
|
|
# Analyze image
|
|
result = classifier.analyze_image("test_image.jpg")
|
|
|
|
print(f"\nDetected objects: {len(result.detected_objects)}")
|
|
print(f"Scene type: {result.scene_type}")
|
|
print(f"Confidence: {result.confidence}")
|
|
print(f"Has embeddings: {result.embeddings is not None}")
|
|
|
|
return result
|
|
|
|
|
|
if __name__ == '__main__':
|
|
example_real_ai()
|