Image Vision

1import os2import sys3import json4import subprocess5import time6import traceback78try:9    import requests10except ImportError:11    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])12    import requests1314NATIVE_IMAGE_VISION_SERVICE_URL = os.getenv(15    "NATIVE_IMAGE_VISION_SERVICE_URL", "http://native-image-vision-service:8104"16)17_EMBLEMA_VERSION = os.getenv("EMBLEMA_VERSION", "dev")18NATIVE_IMAGE_VISION_SERVICE_IMAGE = os.getenv(19    "NATIVE_IMAGE_VISION_SERVICE_IMAGE",20    f"emblema/native-image-vision-service:{_EMBLEMA_VERSION}",21)22HF_CACHE_HOST_PATH = os.getenv("HF_CACHE_HOST_PATH", "/root/.cache/huggingface")23CONTAINER_NAME = "native-image-vision-service"24INPUT_DIR = "/data/input"252627def start_container():28    """Create and start native-image-vision-service, removing any stale container first."""29    subprocess.run(30        ["docker", "rm", "-f", CONTAINER_NAME],31        capture_output=True, text=True32    )3334    hf_token = os.getenv("HUGGINGFACE_TOKEN", "")35    print(f"Creating container {CONTAINER_NAME}...", file=sys.stderr)36    run_cmd = [37        "docker", "run", "-d",38        "--name", CONTAINER_NAME,39        "--network", "emblema",40        "--gpus", "all",41        "-e", "PORT=8104",42        "-e", "DEVICE=cuda",43        "-e", f"HF_TOKEN={hf_token}",44        "-v", f"{HF_CACHE_HOST_PATH}:/root/.cache/huggingface",45        NATIVE_IMAGE_VISION_SERVICE_IMAGE,46    ]47    result = subprocess.run(run_cmd, capture_output=True, text=True)48    if result.returncode != 0:49        print(f"docker run failed (exit {result.returncode}): {result.stderr}", file=sys.stderr)50        raise RuntimeError(f"Failed to start container: {result.stderr}")5152    # Poll health endpoint53    timeout = 18054    interval = 355    elapsed = 056    health_url = f"{NATIVE_IMAGE_VISION_SERVICE_URL}/health"57    while elapsed < timeout:58        try:59            r = requests.get(health_url, timeout=5)60            if r.status_code == 200:61                print(f"Container healthy (waited {elapsed}s).", file=sys.stderr)62                return63        except requests.ConnectionError:64            pass65        time.sleep(interval)66        elapsed += interval6768    raise RuntimeError(f"Container did not become healthy within {timeout}s")697071def stop_container():72    """Remove the container."""73    try:74        subprocess.run(75            ["docker", "rm", "-f", CONTAINER_NAME],76            capture_output=True, text=True, timeout=3077        )78        print(f"Container {CONTAINER_NAME} removed.", file=sys.stderr)79    except Exception as e:80        print(f"Warning: failed to remove container: {e}", file=sys.stderr)818283def main():84    try:85        input_json = sys.stdin.read()86        execution_input = json.loads(input_json)87        inputs = execution_input.get("inputs", {})8889        image = inputs.get("image", "")90        task = inputs.get("task", "more_detailed_caption")9192        if not image:93            raise ValueError("Image is required")9495        valid_tasks = [96            "caption", "detailed_caption", "more_detailed_caption",97            "ocr", "ocr_with_region", "object_detection", "dense_region_caption",98        ]99        if task not in valid_tasks:100            raise ValueError(f"Invalid task '{task}'. Must be one of: {valid_tasks}")101102        local_path = os.path.join(INPUT_DIR, image)103        if not os.path.exists(local_path):104            raise FileNotFoundError(f"Input image not found: {local_path}")105106        # Start the container107        start_container()108109        try:110            # Send image to service111            with open(local_path, "rb") as f:112                resp = requests.post(113                    f"{NATIVE_IMAGE_VISION_SERVICE_URL}/analyze",114                    files={"image": (os.path.basename(local_path), f, "image/png")},115                    data={"task": task},116                    timeout=300,117                )118119            if resp.status_code != 200:120                try:121                    error_detail = resp.json()122                except Exception:123                    error_detail = resp.text124                raise RuntimeError(125                    f"Vision service returned {resp.status_code}: {error_detail}"126                )127128            result_data = resp.json()129            result_text = result_data.get("text", "")130131            print(f"Vision analysis complete: task={task}", file=sys.stderr)132133            # Flat output — keys match OUTPUT_SCHEMA134            output = {135                "text": result_text,136            }137            print(json.dumps(output, indent=2))138139        finally:140            stop_container()141142    except Exception as e:143        error_output = {144            "error": str(e),145            "errorType": type(e).__name__,146            "traceback": traceback.format_exc(),147        }148        print(json.dumps(error_output), file=sys.stderr)149        sys.exit(1)150151152if __name__ == "__main__":153    main()