1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.video.video_nsfw_detect import VideoNsfwDetect
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18    samples = {
19        "video_path": ["/tmp/test_video.mp4", "/tmp/test_video2.mp4"]
20    }
21    ds = daft.from_pydict(samples)
22
23    constructor_kwargs = {
24        "video_src_type": "video_url",
25        "sample_mode": "by_count_uniform",
26        "count_k": 5,
27        "reduce_mode": "avg",
28    }
29    ds = ds.with_column(
30        "safety_score",
31        aihc_udf(
32            VideoNsfwDetect,
33            construct_args=constructor_kwargs,
34            num_cpus=2,
35            concurrency=2,
36            use_gpu=True,
37        )(col("video_path")),
38    )
39    ds.show()