1from __future__ import annotations
2
3import inspect
4import os
5
6import daft
7from daft import col
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.image.image_egoblur import ImageEgoBlur
10
11
12def _get_optional_float(name: str) -> float | None:
13    value = os.getenv(name)
14    return float(value) if value not in {None, ""} else None
15
16
17def _get_bool(name: str, default: bool = False) -> bool:
18    value = os.getenv(name)
19    if value is None:
20        return default
21    return value.lower() in {"1", "true", "yes", "y", "on"}
22
23
24if __name__ == "__main__":
25    if os.getenv("DAFT_RUNNER", "native") == "ray":
26        import ray
27        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
28        daft.set_runner_ray()
29    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
30
31    model_base = os.getenv("MODEL_PATH", "/opt/aihc/models")
32    model_generation = os.getenv("MODEL_GENERATION", "gen1").lower()
33    if model_generation == "gen2" and "model_generation" not in inspect.signature(ImageEgoBlur.__init__).parameters:
34        raise RuntimeError(
35            "MODEL_GENERATION=gen2 requires the gen2-capable ImageEgoBlur implementation. "
36            "The imported ImageEgoBlur is still the old gen1-only version."
37        )
38
39    face_model_path = os.getenv("FACE_MODEL_PATH", os.path.join(model_base, f"ego_blur_face_{model_generation}.jit"))
40    lp_model_path = os.getenv("LP_MODEL_PATH", os.path.join(model_base, f"ego_blur_lp_{model_generation}.jit"))
41    output_dir = os.getenv("OUTPUT_DIR", "")
42
43    samples = {"image_path": ["https://{bucket}.bj.bcebos.com/face.jpg"]}
44    ds = daft.from_pydict(samples)
45
46    constructor_kwargs = {
47        "face_model_path": face_model_path,
48        "lp_model_path": lp_model_path,
49        "model_generation": model_generation,
50        "camera_name": os.getenv("CAMERA_NAME", "camera-rgb"),
51        "output_dir": output_dir,
52        "return_base64": _get_bool("RETURN_BASE64"),
53    }
54    optional_float_kwargs = {
55        "face_score_threshold": _get_optional_float("FACE_SCORE_THRESHOLD"),
56        "lp_score_threshold": _get_optional_float("LP_SCORE_THRESHOLD"),
57        "nms_iou_threshold": _get_optional_float("NMS_IOU_THRESHOLD"),
58    }
59    constructor_kwargs.update({key: value for key, value in optional_float_kwargs.items() if value is not None})
60
61    ds = ds.with_column(
62        "egoblur_result",
63        aihc_udf(
64            ImageEgoBlur,
65            construct_args=constructor_kwargs,
66            num_gpus=0,
67            batch_size=1,
68            concurrency=1,
69        )(col("image_path")),
70    )
71
72    ds.show()
73#╭────────────────────────────────┬────────────────────────────────────────────────────────────────────╮
74#│ image_path                     ┆ egoblur_result                                                    │
75#│ ---                            ┆ ---                                                               │
76#│ String                         ┆ Struct[base64: String, image_path: String, face_bounding_boxes: …] │
77#╞════════════════════════════════╪════════════════════════════════════════════════════════════════════╡
78#│ https://{bucket}.bj.bcebos.com/┆ {base64: None, image_path: None, face_bounding_boxes: [[...]], …}  │
79#╰────────────────────────────────┴────────────────────────────────────────────────────────────────────╯