File size: 3,709 Bytes
24fe247
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
HTTP client for calling remote HF Space microservices (Face Space + Object Space).
Handles cold-start timeouts (HF free tier sleeps after ~5 min inactivity),
retries on 503, and transparent fallback when a space is unavailable.
"""
import asyncio
import io
import os
from typing import Optional

import aiohttp

from src.core.config import FACE_SPACE_URL, OBJECT_SPACE_URL

SPACE_CALL_TIMEOUT = int(os.getenv("SPACE_CALL_TIMEOUT", "120"))  # 2 min — covers cold-start wakeup
SPACE_CALL_RETRIES = int(os.getenv("SPACE_CALL_RETRIES", "2"))

_session: Optional[aiohttp.ClientSession] = None


async def get_session() -> aiohttp.ClientSession:
    global _session
    if _session is None or _session.closed:
        timeout = aiohttp.ClientTimeout(total=SPACE_CALL_TIMEOUT)
        _session = aiohttp.ClientSession(timeout=timeout)
    return _session


async def close_session() -> None:
    global _session
    if _session and not _session.closed:
        await _session.close()
        _session = None


async def _post_image(url: str, image_bytes: bytes, params: dict) -> Optional[dict]:
    """POST image bytes as multipart form to a remote Space endpoint."""
    session = await get_session()
    for attempt in range(SPACE_CALL_RETRIES + 1):
        try:
            data = aiohttp.FormData()
            data.add_field("file", io.BytesIO(image_bytes), filename="image.jpg", content_type="image/jpeg")
            for k, v in params.items():
                data.add_field(k, str(v))
            async with session.post(url, data=data) as resp:
                if resp.status == 200:
                    return await resp.json()
                if resp.status == 503:
                    # Space is waking up — wait with exponential backoff and retry
                    wait = 15 * (attempt + 1)
                    await asyncio.sleep(wait)
                    continue
                # Non-retryable error
                return None
        except asyncio.TimeoutError:
            if attempt < SPACE_CALL_RETRIES:
                await asyncio.sleep(5)
            continue
        except Exception:
            return None
    return None


async def embed_face(image_bytes: bytes, quality_gate: float = 0.35) -> Optional[list]:
    """
    Call the Face Space to detect and embed all faces in the image.
    Returns list of face dicts or None if the space is unavailable.
    Falls back gracefully — caller should treat None as 'no faces found'.
    """
    if not FACE_SPACE_URL:
        return None
    result = await _post_image(
        f"{FACE_SPACE_URL}/embed/face",
        image_bytes,
        {"quality_gate": quality_gate},
    )
    return result.get("faces") if result else None


async def embed_object(image_bytes: bytes) -> Optional[list]:
    """
    Call the Object Space to embed the full image + detected object crops.
    Returns list of object dicts or None if the space is unavailable.
    """
    if not OBJECT_SPACE_URL:
        return None
    result = await _post_image(
        f"{OBJECT_SPACE_URL}/embed/object",
        image_bytes,
        {},
    )
    return result.get("objects") if result else None


async def ping_spaces() -> dict:
    """Health-check all spaces. Returns dict of space_name → bool."""
    session = await get_session()
    results = {}
    for name, url in [("face", FACE_SPACE_URL), ("object", OBJECT_SPACE_URL)]:
        if not url:
            results[name] = None
            continue
        try:
            async with session.get(f"{url}/health", timeout=aiohttp.ClientTimeout(total=5)) as resp:
                results[name] = resp.status == 200
        except Exception:
            results[name] = False
    return results