Models¶

def track(self, img):
    """Detect and track on a BGR frame. Returns [Tracks][vizor.boxes.Tracks]."""
    raise NotImplementedError(f"{type(self).__name__} cannot be a primary model")

def find(self, img, names=None):
    """Detect on a whole BGR frame. Returns [Preds][vizor.boxes.Preds]."""
    raise NotImplementedError(f"{type(self).__name__} does not support mode='full'")

def name(self, crop, names=None, hint=None):
    """Classify a BGR crop. Returns a class id, or None if unsure.

    ``hint`` is what the primary thought the object was.
    """
    raise NotImplementedError(f"{type(self).__name__} does not support mode='crop'")

def batch(self, crops, names=None, hints=None):
    """Classify several BGR crops. Returns one class id or None per crop.

    The default asks ``name`` once per crop, so a model that only implements
    ``name`` works unchanged. Override this when the model can do the whole
    list in one call, which is what makes crop mode cheap.

    ``hints`` is what the primary thought each object was, in the same order.
    """
    hints = list(hints) if hints is not None else [None] * len(crops)
    return [self.name(c, names, h) for c, h in zip(crops, hints)]

def grid(self, collages, names=None, hints=None, tiles=1):
    """Classify collages, each one object shown ``tiles`` times over a video.

    Returns one class id or None per collage, the same shape as ``batch``.
    The default forwards to ``batch``, which treats a collage as an ordinary
    image, so a model that has never heard of collages still answers.
    Override it to word the prompt for a grid.
    """
    return self.batch(collages, names, hints)

def reset(self):
    """Drop any per-video state. Called by ``Vizor.reset``."""

@staticmethod
def encode(img, quality=90):
    """BGR array to a base64 data url the chat APIs accept."""
    import base64

    ok, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        raise ValueError("could not encode crop as jpeg")
    return "data:image/jpeg;base64," + base64.b64encode(buf).decode()

def ask(self, imgs, text, **kw):
    """Send one or more images with one question, return the reply as a string.

    Several images go in a single message, each preceded by its number, so
    the model can answer about all of them at once.
    """
    if not isinstance(imgs, (list, tuple)):
        imgs = [imgs]
    content = []
    for i, img in enumerate(imgs, 1):
        if len(imgs) > 1:
            content.append({"type": "text", "text": f"Crop {i}:"})
        content.append({"type": "image_url", "image_url": {"url": self.encode(img)}})
    content.append({"type": "text", "text": text})
    reply = self.client.chat.completions.create(
        model=self.model,
        messages=[{"role": "user", "content": content}],
        **{**self.kw, **kw},
    )
    return reply.choices[0].message.content

def name(self, crop, names=None, hint=None):
    """Ask the model which class the crop is. Returns a class id, or None if unsure."""
    names = names if names is not None else self.names
    text = self.prompt.format(hint=hint, menu=menu(names))
    return parse_id(self.ask(crop, text), ids(names))

def batch(self, crops, names=None, hints=None):
    """Classify several crops, ``chunk`` of them per request.

    One request carrying eight crops costs one round trip instead of eight,
    which is most of the wall clock in crop mode. The trade is that the model
    has to keep the order straight, and a reply with the wrong count loses
    the crops it did not cover rather than shifting the rest.
    """
    names = names if names is not None else self.names
    hints = list(hints) if hints is not None else [None] * len(crops)
    valid, lines = ids(names), menu(names)
    out = []
    for i in range(0, len(crops), self.chunk):
        part, hint = crops[i:i + self.chunk], hints[i:i + self.chunk]
        if len(part) == 1:
            out.append(self.name(part[0], names, hint[0]))
            continue
        text = self.batch_prompt.format(
            n=len(part),
            hints=", ".join(f"{j}. {h!r}" for j, h in enumerate(hint, 1)),
            menu=lines,
        )
        # 16 tokens holds one id, not eight, so give the reply room to fit
        room = max(self.kw.get("max_tokens", 16), 8 * len(part))
        out += parse_ids(self.ask(part, text, max_tokens=room), len(part), valid)
    return out

def grid(self, collages, names=None, hints=None, tiles=1):
    """Ask about one collage per request. Returns a class id or None per collage.

    ``chunk`` does not apply here. Several grids in one request would ask the
    model to keep the tiles and the grids straight at the same time, and
    collage mode already fires once per track rather than once per frame, so
    the round trips are not where the cost is.
    """
    names = names if names is not None else self.names
    hints = list(hints) if hints is not None else [None] * len(collages)
    valid, lines = ids(names), menu(names)
    out = []
    for sheet, hint in zip(collages, hints):
        text = self.grid_prompt.format(n=tiles, hint=hint, menu=lines)
        out.append(parse_id(self.ask(sheet, text), valid))
    return out

def ask(self, img, question):
    """Send one image and one question, return the decoded reply."""
    inputs = self.processor(text=self._text(question), images=[_pil(img)],
                            return_tensors="pt").to(self.device, self.dtype)
    out = self.model.generate(**inputs, **self.gen)
    out = out[:, inputs["input_ids"].shape[-1]:]  # drop the echoed prompt
    return self.processor.batch_decode(out, skip_special_tokens=True)[0].strip()

def name(self, crop, names=None, hint=None):
    """Ask the model which class the crop is. Returns a class id, or None if unsure."""
    names = names if names is not None else self.names
    return parse_id(self.ask(crop, self.prompt.format(hint=hint, menu=menu(names))),
                    ids(names))

def grid(self, collages, names=None, hints=None, tiles=1):
    """Ask about each collage in its own forward pass, with the grid prompt."""
    names = names if names is not None else self.names
    hints = list(hints) if hints is not None else [None] * len(collages)
    valid, lines = ids(names), menu(names)
    return [parse_id(self.ask(s, self.grid_prompt.format(n=tiles, hint=h, menu=lines)), valid)
            for s, h in zip(collages, hints)]

def run(self, img, task, text=""):
    """Run one Florence task and return its parsed dict."""
    pil = _pil(img)
    inputs = self.processor(text=task + text, images=pil, return_tensors="pt")
    inputs = inputs.to(self.device, self.dtype)
    out = self.model.generate(input_ids=inputs["input_ids"],
                              pixel_values=inputs["pixel_values"], **self.gen)
    raw = self.processor.batch_decode(out, skip_special_tokens=False)[0]
    return self.processor.post_process_generation(raw, task=task, image_size=pil.size)

def find(self, img, names=None):
    """Ground the class names in a BGR frame. Every box comes back at confidence 1.0."""
    lookup = self._lookup(names)
    task = self.task
    text = ", ".join(lookup) if task == "<CAPTION_TO_PHRASE_GROUNDING>" and lookup else ""
    out = self.run(img, task, text).get(task, {})
    rows = []
    for box, label in zip(out.get("bboxes", []), out.get("labels", [])):
        cls = lookup.get(str(label).lower().strip())
        if cls is None:
            continue
        rows.append([*box, 1.0, cls, -1])
    data = np.array(rows, np.float32) if rows else np.zeros((0, 7), np.float32)
    return Preds(data, names=names if names is not None else self.names)

def name(self, crop, names=None, hint=None):
    """Class of the largest thing Florence finds in the crop."""
    preds = self.find(crop, names)
    if not len(preds):
        return None
    areas = (preds.boxes[:, 2] - preds.boxes[:, 0]) * (preds.boxes[:, 3] - preds.boxes[:, 1])
    return int(preds.cls[int(areas.argmax())])

def grid(self, collages, names=None, hints=None, tiles=1):
    """Largest thing Florence grounds in each collage. It has no chat prompt to word."""
    return [self.name(sheet, names) for sheet in collages]

def reset(self):
    """Rewind to the first frame."""
    self.i = 0

def next(self):
    """The next frame as a float32 array, or an empty one past the end."""
    if self.i >= len(self.frames):
        return np.zeros((0, 7), np.float32)
    data = np.asarray(self.frames[self.i], dtype=np.float32)
    self.i += 1
    if data.ndim == 1:
        data = data.reshape(0, 7) if not data.size else data.reshape(1, -1)
    if self.cols is not None and len(data):
        data = data[:, self.cols]
    return data

def track(self, img=None):
    """The next saved frame as Tracks. ``img`` is ignored."""
    return Tracks(self.next(), names=self.names)

def find(self, img=None, names=None):
    """The next saved frame as Preds. ``img`` is ignored."""
    return Preds(self.next(), names=names or self.names)