miku-discord/bot/utils/face_detector_manager.py

# face_detector_manager.py
"""
Manages on-demand starting/stopping of anime-face-detector container
to free up VRAM when not needed.
"""

import asyncio
import aiohttp
import subprocess
import time
from typing import Optional, Dict


class FaceDetectorManager:
    """Manages the anime-face-detector container lifecycle"""
    
    FACE_DETECTOR_API = "http://anime-face-detector:6078/detect"
    HEALTH_ENDPOINT = "http://anime-face-detector:6078/health"
    CONTAINER_NAME = "anime-face-detector"
    STARTUP_TIMEOUT = 30  # seconds
    
    def __init__(self):
        self.is_running = False
    
    async def start_container(self, debug: bool = False) -> bool:
        """
        Start the anime-face-detector container.
        
        Returns:
            True if started successfully, False otherwise
        """
        try:
            if debug:
                print("🚀 Starting anime-face-detector container...")
            
            # Start container using docker compose
            result = subprocess.run(
                ["docker", "compose", "up", "-d", self.CONTAINER_NAME],
                cwd="/app",  # Assumes we're in the bot container, adjust path as needed
                capture_output=True,
                text=True,
                timeout=30
            )
            
            if result.returncode != 0:
                if debug:
                    print(f"⚠️ Failed to start container: {result.stderr}")
                return False
            
            # Wait for API to be ready
            start_time = time.time()
            while time.time() - start_time < self.STARTUP_TIMEOUT:
                if await self._check_health():
                    self.is_running = True
                    if debug:
                        print(f"✅ Face detector container started and ready")
                    return True
                await asyncio.sleep(1)
            
            if debug:
                print(f"⚠️ Container started but API not ready after {self.STARTUP_TIMEOUT}s")
            return False
            
        except Exception as e:
            if debug:
                print(f"⚠️ Error starting face detector container: {e}")
            return False
    
    async def stop_container(self, debug: bool = False) -> bool:
        """
        Stop the anime-face-detector container to free VRAM.
        
        Returns:
            True if stopped successfully, False otherwise
        """
        try:
            if debug:
                print("🛑 Stopping anime-face-detector container...")
            
            result = subprocess.run(
                ["docker", "compose", "stop", self.CONTAINER_NAME],
                cwd="/app",
                capture_output=True,
                text=True,
                timeout=15
            )
            
            if result.returncode == 0:
                self.is_running = False
                if debug:
                    print("✅ Face detector container stopped")
                return True
            else:
                if debug:
                    print(f"⚠️ Failed to stop container: {result.stderr}")
                return False
                
        except Exception as e:
            if debug:
                print(f"⚠️ Error stopping face detector container: {e}")
            return False
    
    async def _check_health(self) -> bool:
        """Check if the face detector API is responding"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    self.HEALTH_ENDPOINT,
                    timeout=aiohttp.ClientTimeout(total=2)
                ) as response:
                    return response.status == 200
        except:
            return False
    
    async def detect_face_with_management(
        self,
        image_bytes: bytes,
        unload_vision_model: callable = None,
        reload_vision_model: callable = None,
        debug: bool = False
    ) -> Optional[Dict]:
        """
        Detect face with automatic container lifecycle management.
        
        Args:
            image_bytes: Image data as bytes
            unload_vision_model: Optional callback to unload vision model first
            reload_vision_model: Optional callback to reload vision model after
            debug: Enable debug output
        
        Returns:
            Detection dict or None
        """
        container_was_started = False
        
        try:
            # Step 1: Unload vision model if callback provided
            if unload_vision_model:
                if debug:
                    print("📤 Unloading vision model to free VRAM...")
                await unload_vision_model()
                await asyncio.sleep(2)  # Give time for VRAM to clear
            
            # Step 2: Start face detector if not running
            if not self.is_running:
                if not await self.start_container(debug=debug):
                    if debug:
                        print("⚠️ Could not start face detector container")
                    return None
                container_was_started = True
            
            # Step 3: Detect face
            result = await self._detect_face_api(image_bytes, debug=debug)
            
            return result
            
        finally:
            # Step 4: Stop container and reload vision model
            if container_was_started:
                await self.stop_container(debug=debug)
            
            if reload_vision_model:
                if debug:
                    print("📥 Reloading vision model...")
                await reload_vision_model()
    
    async def _detect_face_api(self, image_bytes: bytes, debug: bool = False) -> Optional[Dict]:
        """Call the face detection API"""
        try:
            async with aiohttp.ClientSession() as session:
                form = aiohttp.FormData()
                form.add_field('file', image_bytes, filename='image.jpg', content_type='image/jpeg')
                
                async with session.post(
                    self.FACE_DETECTOR_API,
                    data=form,
                    timeout=aiohttp.ClientTimeout(total=30)
                ) as response:
                    if response.status != 200:
                        if debug:
                            print(f"⚠️ Face detection API returned status {response.status}")
                        return None
                    
                    result = await response.json()
                    
                    if result.get('count', 0) == 0:
                        if debug:
                            print("👤 No faces detected by API")
                        return None
                    
                    detections = result.get('detections', [])
                    if not detections:
                        return None
                    
                    best_detection = max(detections, key=lambda d: d.get('confidence', 0))
                    bbox = best_detection.get('bbox', [])
                    confidence = best_detection.get('confidence', 0)
                    keypoints = best_detection.get('keypoints', [])
                    
                    if len(bbox) >= 4:
                        x1, y1, x2, y2 = bbox[:4]
                        center_x = int((x1 + x2) / 2)
                        center_y = int((y1 + y2) / 2)
                        
                        if debug:
                            width = int(x2 - x1)
                            height = int(y2 - y1)
                            print(f"👤 Detected {len(detections)} face(s) via API, using best at ({center_x}, {center_y}) [confidence: {confidence:.2%}]")
                            print(f"   Bounding box: x={int(x1)}, y={int(y1)}, w={width}, h={height}")
                            print(f"   Keypoints: {len(keypoints)} facial landmarks detected")
                        
                        return {
                            'center': (center_x, center_y),
                            'bbox': bbox,
                            'confidence': confidence,
                            'keypoints': keypoints,
                            'count': len(detections)
                        }
        
        except Exception as e:
            if debug:
                print(f"⚠️ Error calling face detection API: {e}")
        
        return None


# Global instance
face_detector_manager = FaceDetectorManager()