Florence-2: A Unified Model for Many Vision Tasks

# A supported subset of the release processor's task dispatch.
TASKS = {
    "<CAPTION>": ("text", "What does the image describe?"),
    "<OD>": ("boxes", "Locate the objects with category name in the image."),
    "<DENSE_REGION_CAPTION>": (
        "boxes",
        "Locate the objects in the image, with their descriptions.",
    ),
    "<OCR>": ("text", "What is the text in the image?"),
    "<OCR_WITH_REGION>": (
        "ocr_quad",
        "What is the text in the image, with regions?",
    ),
    "<REFERRING_EXPRESSION_SEGMENTATION>": (
        "polygon",
        "Locate {input} in the image with mask",
    ),
}
GRAMMARS = {
    "boxes": "label + four location tokens (x0,y0,x1,y1), repeated",
    "text": "plain text, without spatial tokens",
    "ocr_quad": "text + eight location tokens (four clockwise vertices)",
    "polygon": "clockwise x,y pairs; <sep> parts and <poly>...</poly> groups",
}


def task_prompt(task, text=""):
    """Dispatch a supported task and construct its released text prompt.

    Args:
        task: Exact task marker from TASKS.
        text: Referring expression for the segmentation task only.

    Returns:
        A dictionary with task, prompt, format, and declared grammar.

    Raises:
        ValueError: For unsupported tasks or incorrect task input.
    """
    if not isinstance(task, str) or task not in TASKS:
        raise ValueError("Unsupported task marker.")
    if not isinstance(text, str):
        raise ValueError("Task input must be text.")
    kind, template = TASKS[task]
    needs_input = "{input}" in template
    if needs_input and (not text.strip() or "<" in text or ">" in text):
        raise ValueError("A nonempty referring expression is required.")
    if not needs_input and text:
        raise ValueError("This task takes no additional text.")
    return {
        "task": task,
        "prompt": template.format(input=text),
        "format": kind,
        "grammar": GRAMMARS[kind],
    }

def quantize(boxes, size):
    """Floor and clamp [N,4] boxes using width for x and height for y.

    Coordinates outside the image are clamped, matching BoxQuantizer. The
    half-bin error bound applies to coordinates inside the closed image.
    """
    scales = image_scale(size)
    values = box_tensor(boxes)
    return (values / (scales / BINS)).floor().clamp(0, BINS - 1).long()

def serialize(labels, boxes, size, task="<OD>"):
    """Serialize annotated boxes, not inferred boxes, as label + four tokens.

    This strict subset requires nonempty labels without token delimiters.
    Polygon parts and OCR quadrilaterals have separate grammars and are not
    accepted by this four-coordinate serializer. No BPE tokenizer is modeled.
    """
    if task_prompt(task)["format"] != "boxes":
        raise ValueError("The box serializer requires a box task.")
    bins = quantize(boxes, size)
    if not isinstance(labels, (list, tuple)) or len(labels) != len(bins):
        raise ValueError("One label is required per box.")
    pieces = []
    for label, corners in zip(labels, bins.tolist()):
        if (
            not isinstance(label, str)
            or not label.strip()
            or "<" in label
            or ">" in label
            or label != label.strip()
        ):
            raise ValueError("Labels must be stripped text without brackets.")
        pieces.append(label + "".join(f"<loc_{bin_id}>" for bin_id in corners))
    return "".join(pieces)

def teacher_forcing(target_ids, logits):
    """Right-shift decoder inputs once; score same-position target labels.

    Args:
        target_ids: CPU int64 [B,T] IDs, with PAD=1 only in suffix padding.
        logits: Finite CPU floating [B,T,V] fixed prediction logits.

    Returns:
        Decoder input IDs, labels with PAD replaced by -100, and mean CE.
        All-padding batches return differentiable zero (authored extension).

    Raises:
        ValueError: For invalid shapes, types, IDs, or padding placement.
    """
    if (
        not isinstance(target_ids, torch.Tensor)
        or target_ids.dtype != torch.long
        or target_ids.ndim != 2
        or min(target_ids.shape) < 1
        or target_ids.device.type != "cpu"
    ):
        raise ValueError("Targets must be nonempty CPU int64 [B,T].")
    if (
        not isinstance(logits, torch.Tensor)
        or logits.ndim != 3
        or logits.shape[:2] != target_ids.shape
        or logits.shape[2] <= 2
        or not logits.is_floating_point()
        or logits.device.type != "cpu"
        or not torch.isfinite(logits).all()
    ):
        raise ValueError("Logits must be finite CPU float [B,T,V].")
    if (target_ids < 0).any() or (target_ids >= logits.shape[2]).any():
        raise ValueError("Target IDs must belong to the logit vocabulary.")
    padding = target_ids == PAD
    if (padding[:, :-1] & ~padding[:, 1:]).any():
        raise ValueError("Padding must be a suffix.")
    labels = target_ids.masked_fill(padding, -100)
    decoder_inputs = torch.full_like(target_ids, PAD)
    decoder_inputs[:, 0] = DECODER_START
    decoder_inputs[:, 1:] = labels[:, :-1]
    decoder_inputs.masked_fill_(decoder_inputs == -100, PAD)
    if (~padding).any():
        loss = functional.cross_entropy(
            logits.reshape(-1, logits.shape[-1]),
            labels.reshape(-1),
            ignore_index=-100,
            reduction="mean",
        )
    else:
        loss = logits.sum() * 0
    return decoder_inputs, labels, loss

def dequantize(bins, size):
    """Decode integer [N,4] bins to centers in original image coordinates."""
    scales = image_scale(size)
    values = torch.as_tensor(bins)
    if (
        values.device.type != "cpu"
        or values.ndim != 2
        or values.shape[1] != 4
        or values.dtype == torch.bool
        or values.is_floating_point()
        or values.is_complex()
        or (values < 0).any()
        or (values >= BINS).any()
    ):
        raise ValueError("Location bins must be integer [N,4] in [0,999].")
    decoded = (values.to(torch.float64) + 0.5) * (scales / BINS)
    return box_tensor(decoded)


def parse_dequantize(sequence, task, size):
    """Parse the complete box/text body; reject every malformed remainder.

    Unlike the release regex parser, this authored strict interface never
    skips unmatched spans or accepts out-of-range bins. Caller supplies the
    body only, without BOS/EOS/PAD wrappers. Polygon/OCR-region formats are
    declared but deliberately unsupported, rather than read as boxes.
    """
    kind = task_prompt(
        task, "object" if task == "<REFERRING_EXPRESSION_SEGMENTATION>" else ""
    )["format"]
    image_scale(size)
    if not isinstance(sequence, str):
        raise ValueError("Sequence body must be text.")
    if kind in ("polygon", "ocr_quad"):
        raise NotImplementedError("This interface only parses boxes/text.")
    if kind == "text":
        if "<" in sequence or ">" in sequence:
            raise ValueError("Text output cannot contain special tokens.")
        return {task: sequence}
    record = re.compile(r"([^<>]+)((?:<loc_[0-9]+>){4})")
    labels, bins, cursor = [], [], 0
    while cursor < len(sequence):
        match = record.match(sequence, cursor)
        if match is None or match[1] != match[1].strip():
            raise ValueError("Malformed label/box sequence.")
        labels.append(match[1])
        bins.append([int(value) for value in re.findall(r"[0-9]+", match[2])])
        cursor = match.end()
    corners = torch.tensor(bins, dtype=torch.long).reshape(-1, 4)
    boxes = dequantize(corners, size).tolist()
    return {task: {"bboxes": boxes, "labels": labels}}