import json import time import cv2 import numpy as np import google.generativeai as genai from pathlib import Path from typing import Optional from uuid import UUID from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from Backend.core.config import settings from Backend.core.events import event_bus, IssueClassified, IssueCreated from Backend.core.logging import get_logger from Backend.core.schemas import ClassificationResult, DetectionBox, CLASS_ID_TO_CATEGORY, IssueCategory from Backend.database.models import Classification, Issue, IssueImage, IssueEvent from Backend.orchestration.base import BaseAgent from Backend.utils.fuzzy_match import auto_validate_issue from Backend.utils.storage import save_bytes, download_from_supabase, get_upload_url logger = get_logger(__name__, agent_name="VisionAgent") if settings.gemini_api_key: genai.configure(api_key=settings.gemini_api_key) class VisionAgent(BaseAgent): _model = None def __init__(self, db: Optional[AsyncSession] = None): super().__init__("VisionAgent") self.db = db if settings.gemini_api_key: self.gemini_model = genai.GenerativeModel('gemma-3-27b-it') else: self.gemini_model = None @classmethod def load_model(cls): if cls._model is None: from ultralytics import YOLO model_path = settings.model_path if not model_path.exists(): raise FileNotFoundError(f"Model not found: {model_path}") cls._model = YOLO(str(model_path)) logger.info(f"YOLO model loaded from {model_path}") return cls._model @classmethod def get_model(cls): if cls._model is None: cls.load_model() return cls._model async def download_image(self, remote_path: str) -> bytes: return await download_from_supabase(remote_path) async def save_annotated( self, image_data: bytes, detections: list[DetectionBox], primary_class_id: Optional[int], original_path: str, subfolder: str, ) -> str: nparr = np.frombuffer(image_data, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if img is None: raise ValueError("Invalid image data") if primary_class_id is not None: draw_detections = [d for d in detections if d.class_id == primary_class_id] else: draw_detections = [] for d in draw_detections: x1, y1, x2, y2 = map(int, d.bbox) cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2) label = f"{d.class_name} {d.confidence:.2f}" cv2.putText(img, label, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3) original_name = Path(original_path).stem annotated_filename = f"annotated_{original_name}.jpg" _, buffer = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 90]) image_bytes = buffer.tobytes() remote_path = await save_bytes(image_bytes, annotated_filename, subfolder=subfolder) return remote_path async def run_inference(self, image_data: bytes) -> tuple[list, float]: model = self.get_model() nparr = np.frombuffer(image_data, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if img is None: raise ValueError("Invalid image data") start_time = time.perf_counter() results = model.predict( source=img, conf=settings.model_confidence_threshold, imgsz=settings.model_input_size, verbose=False, ) inference_time = (time.perf_counter() - start_time) * 1000 return results, inference_time async def gemini_classify_image( self, image_data: bytes, description: Optional[str] = None ) -> tuple[Optional[IssueCategory], float, Optional[str]]: if not self.gemini_model: return None, 0.0, None allowed = [ {"class_id": k, "class_name": v.value} for k, v in CLASS_ID_TO_CATEGORY.items() ] prompt = ( "Classify the photo into exactly one of the allowed categories. " "Return ONLY valid JSON with keys: class_id (int), confidence (0.0-1.0), reasoning (max 80 chars).\n\n" f"Allowed categories: {json.dumps(allowed)}\n" f"User description: {(description or '')[:200]}" ) try: response = self.gemini_model.generate_content( [ {"text": prompt}, { "inline_data": { "mime_type": "image/jpeg", "data": image_data, } }, ] ) text = (response.text or "").replace("```json", "").replace("```", "").strip() data = json.loads(text) class_id = data.get("class_id") confidence = float(data.get("confidence", 0.0)) reasoning = data.get("reasoning") if not isinstance(class_id, int): return None, 0.0, None category = CLASS_ID_TO_CATEGORY.get(class_id) if not category: return None, 0.0, None confidence = max(0.0, min(1.0, confidence)) return category, confidence, reasoning except Exception as e: logger.error(f"Gemini vision classification failed: {e}") return None, 0.0, None def extract_detections(self, results) -> list[DetectionBox]: detections = [] for result in results: boxes = result.boxes if boxes is not None: for i in range(len(boxes)): class_id = int(boxes.cls[i].item()) confidence = float(boxes.conf[i].item()) bbox = tuple(boxes.xyxy[i].tolist()) category = CLASS_ID_TO_CATEGORY.get(class_id) if category: detections.append(DetectionBox( class_id=class_id, class_name=category.value, confidence=confidence, bbox=bbox, )) return detections async def classify_image( self, image_path: str, subfolder: str = "", description: Optional[str] = None ) -> tuple[list[DetectionBox], str, Optional[IssueCategory], float, Optional[str]]: image_data = await self.download_image(image_path) results, inference_time = await self.run_inference(image_data) detections = self.extract_detections(results) primary_class_id = max(detections, key=lambda d: d.confidence).class_id if detections else None annotated_path = await self.save_annotated(image_data, detections, primary_class_id, image_path, subfolder) gemini_category = None gemini_confidence = 0.0 gemini_reasoning = None if self.gemini_model and detections and max(d.confidence for d in detections) < 0.5: gemini_category, gemini_confidence, gemini_reasoning = await self.gemini_classify_image( image_data=image_data, description=description ) logger.info(f"Inference completed in {inference_time:.2f}ms, {len(detections)} detections") return detections, annotated_path, gemini_category, gemini_confidence, gemini_reasoning async def process_issue( self, issue_id: UUID, image_paths: list[str], description: Optional[str] = None ) -> ClassificationResult: all_detections = [] annotated_paths = [] total_time = 0.0 subfolder = str(issue_id) gemini_best_category = None gemini_best_confidence = 0.0 gemini_best_reasoning = None for path in image_paths: start = time.perf_counter() detections, annotated_path, gemini_category, gemini_confidence, gemini_reasoning = await self.classify_image( path, subfolder=subfolder, description=description ) total_time += (time.perf_counter() - start) * 1000 all_detections.extend(detections) annotated_paths.append(annotated_path) if gemini_category and gemini_confidence > gemini_best_confidence: gemini_best_category = gemini_category gemini_best_confidence = gemini_confidence gemini_best_reasoning = gemini_reasoning if self.db: query = select(IssueImage).where(IssueImage.file_path == path) result = await self.db.execute(query) image_record = result.scalar_one_or_none() if image_record: image_record.annotated_path = annotated_path result = ClassificationResult( issue_id=issue_id, detections=all_detections, annotated_urls=[get_upload_url(p) for p in annotated_paths], inference_time_ms=total_time, ) if gemini_best_category and (not result.primary_category or result.primary_confidence < 0.5): result.primary_category = gemini_best_category result.primary_confidence = gemini_best_confidence detected_categories = list(set(d.class_name for d in all_detections)) auto_validated, validation_reason = auto_validate_issue(description, detected_categories) validation_source = "auto" if auto_validated else "pending_manual" new_state = "validated" if auto_validated else "reported" self.log_decision( issue_id=issue_id, decision=f"Validation: {validation_source}", reasoning=validation_reason ) if self.db: classification = Classification( issue_id=issue_id, primary_category=result.primary_category.value if result.primary_category else None, primary_confidence=result.primary_confidence, detections_json=json.dumps([d.model_dump() for d in all_detections]), inference_time_ms=total_time, ) self.db.add(classification) issue = await self.db.get(Issue, issue_id) if issue: issue.state = new_state issue.validation_source = validation_source issue.validation_reason = validation_reason event_record = IssueEvent( issue_id=issue_id, event_type="classified", agent_name=self.name, event_data=json.dumps({ "category": result.primary_category.value if result.primary_category else None, "confidence": result.primary_confidence, "detections_count": len(all_detections), "validation_source": validation_source, "annotated_images": annotated_paths, "gemini_category": gemini_best_category.value if gemini_best_category else None, "gemini_confidence": gemini_best_confidence, "gemini_reasoning": gemini_best_reasoning, }) ) self.db.add(event_record) await self.db.flush() if result.primary_category: event = IssueClassified( issue_id=issue_id, category=result.primary_category.value, confidence=result.primary_confidence, detections_count=len(all_detections), metadata={ "validation_source": validation_source, "validation_reason": validation_reason, "annotated_images": [get_upload_url(p) for p in annotated_paths], } ) await event_bus.publish(event) return result async def handle(self, event: IssueCreated) -> None: await self.process_issue( event.issue_id, event.image_paths, event.description )