src/prompts.py

"""Prompt templates for space debris localization tasks."""

from dataclasses import dataclass


@dataclass(frozen=True)
class PromptTemplate:
    """A reusable prompt template with metadata."""

    name: str
    template: str
    description: str
    category: str


DETECTION_TEMPLATES: list[PromptTemplate] = [
    PromptTemplate(
        name="debris_single",
        template="Locate a single instance that matches the following description: {phrase}.",
        description="Locate one instance of a specific object",
        category="grounding",
    ),
    PromptTemplate(
        name="debris_multi",
        template="Locate all the instances that match the following description: {phrase}.",
        description="Locate all instances of a specific object type",
        category="detection",
    ),
    PromptTemplate(
        name="debris_categories",
        template="Locate all the instances that matches the following description: {categories}.",
        description="Detect multiple object categories at once",
        category="detection",
    ),
    PromptTemplate(
        name="text_grounding",
        template="Please locate the text referred as {phrase}.",
        description="Locate text labels or markings in the image",
        category="text",
    ),
    PromptTemplate(
        name="scene_text",
        template="Detect all the text in box format.",
        description="Detect all visible text in the scene",
        category="text",
    ),
]

SPACE_DEBRIS_EXAMPLES: list[dict[str, str]] = [
    {
        "phrase": "space debris",
        "prompt": "Locate all the instances that match the following description: space debris.",
        "description": "Find all visible space debris fragments",
    },
    {
        "phrase": "satellite fragment",
        "prompt": "Locate all the instances that match the following description: satellite fragment.",
        "description": "Identify broken satellite pieces",
    },
    {
        "phrase": "solar panel",
        "prompt": "Locate all the instances that match the following description: solar panel.",
        "description": "Find satellite solar panels",
    },
    {
        "phrase": "antenna",
        "prompt": "Locate all the instances that match the following description: antenna.",
        "description": "Locate spacecraft antennas",
    },
    {
        "phrase": "rocket body",
        "prompt": "Locate all the instances that match the following description: rocket body.",
        "description": "Find spent rocket stages",
    },
    {
        "phrase": "spacecraft",
        "prompt": "Locate a single instance that matches the following description: spacecraft.",
        "description": "Locate a single spacecraft",
    },
    {
        "phrase": "debris field",
        "prompt": "Locate all the instances that match the following description: debris field.",
        "description": "Find clusters of orbital debris",
    },
    {
        "phrase": "thermal blanket",
        "prompt": "Locate all the instances that match the following description: thermal blanket.",
        "description": "Find loose thermal insulation material",
    },
]


def build_detect_prompt(categories: list[str]) -> str:
    """Build a multi-category detection prompt."""
    joined = "</c>".join(categories)
    return f"Locate all the instances that matches the following description: {joined}."


def build_grounding_prompt(phrase: str, *, single: bool = False) -> str:
    """Build a phrase grounding prompt."""
    if single:
        return f"Locate a single instance that matches the following description: {phrase}."
    return f"Locate all the instances that match the following description: {phrase}."


def get_example_prompts() -> list[list[str]]:
    """Return example prompts for Gradio examples component."""
    return [[ex["prompt"]] for ex in SPACE_DEBRIS_EXAMPLES]