Dataset Studio quickstart

import os
from getpass import getpass
from typing import Literal

from nomadic import NomadicAI
from pydantic import BaseModel, Field

# Uses NOMADIC_API_KEY if set; otherwise asks for your key.
client = NomadicAI(
    api_key=os.environ.get("NOMADIC_API_KEY") or getpass("API key: ")
)

# 1. Create a Studio from a public batch of 75 robot-arm episodes.
studio = client.studio.create_from_batch(
    "aa1dbff79fb64c2a9f8a6bb92d322894",
    title="Robot actions quickstart",
)
print("Studio:", studio.url)

# 2. Label your events: was the object on a surface or in a container?
class PickLocation(BaseModel):
    pick_from: Literal["surface", "container"] = Field(
        description=(
            "Where the robot picks the object up from: "
            "an open surface such as a table, desk, or counter, "
            "or inside a container such as a cup, mug, bowl, "
            "drawer, cabinet, or sink"
        )
    )

client.studio.add_taxonomy(studio.id, PickLocation).wait()

# 3. Search in natural language, and play the matches.
matches = client.studio.search(
    studio.id, "robot lifts a pen out of a mug", top_k=5
)
matches.visualize(max_videos=3)

# 4. Filter to the label values you choose.
events = client.studio.filter(
    studio.id,
    filters={
        "object_category": "writing_tool",
        "pick_from": "container",
    },
)
print("Matching events:", events.total)

# 5. Export them (paid, sponsored, or enterprise plans).
selection = client.studio.select(studio.id).include(events)
job = selection.export()
job.wait()
files = job.download("./studio-export")
print("Downloaded to:", files["output_dir"])