Models¶

@abstractmethod
def _predict(self, images: list[NumpyImage], **kwargs) -> list[Result]:
    """Model specific prediction method"""

def predict(
    self,
    images: Collection[NumpyImage],
    batch_size: int = 1,
    image_scaling_factor: float = 1.0,
    tqdm_kwargs: dict[str, Any] | None = None,
    **kwargs,
) -> list[Result]:
    """Perform inference on images

    Takes an arbitrary number of inputs and runs batched inference.
    The inputs can be streamed from an iterator and don't need to
    be simultaneously read into memory. Prints a progress bar using
    `tqdm`. This is a template method which uses the model-specific
    `_predict(...)`.

    Arguments:
        images: Input images
        batch_size: Inference batch size, defaults to 1
        image_scaling_factor: If < 1, all input images will be down-
            scaled by this factor, which can be useful for speeding
            up inference on higher resolution images. All geometric
            data in the result (e.g., bounding boxes) are reported
            with respect to the original resolution.
        tqdm_kwargs: Optional keyword arguments to control the
            progress bar.
        **kwargs: Optional keyword arguments that are forwarded to
            the model specific prediction method `_predict(...)`.
    """

    batch_size = max(batch_size, 1)  # make sure batch size is at least 1
    image_scaling_factor = max(10e-10, min(image_scaling_factor, 1))  # clip scaling factor to (0, 1]

    n_batches = (len(images) + batch_size - 1) // batch_size
    model_name = self.__class__.__name__
    logger.info(
        "Model '%s' on device '%s' received %d images in batches of %d images per batch (%d batches)",
        model_name,
        self.device,
        len(images),
        batch_size,
        n_batches,
    )

    results = []
    batches = _batch(images, batch_size)
    desc = f"{model_name}: Running inference (batch size {batch_size})"
    for i, batch in enumerate(tqdm(batches, desc, n_batches, **(tqdm_kwargs or {}))):
        msg = "%s: Running inference on %d images (batch %d of %d)"
        logger.info(msg, model_name, len(batch), i + 1, n_batches)
        scaled_batch = [rescale_linear(image, image_scaling_factor) for image in batch]
        batch_results = self._predict(scaled_batch, **kwargs)
        for result in batch_results:
            result.rescale(1 / image_scaling_factor)
            results.append(result)
    return results

def _predict(self, images: list[np.ndarray], **generation_kwargs) -> list[Result]:
    """TrOCR-specific prediction method.

    This method is used by `predict()` and should typically not be
    called directly. However, `predict()` forwards additional kwargs
    to this method.

    Arguments:
        images: Input images.
        **generation_kwargs: Optional keyword arguments that are
            forwarded to the model's .generate() method.
    """

    # Prepare generation keyword arguments: Generally, all kwargs are
    # forwarded to the model's .generate method, but some need to be
    # explicitly set (and possibly overridden) to ensure that we get the
    # output format we want.
    generation_kwargs["num_return_sequences"] = generation_kwargs.get("num_beams", 1)
    generation_kwargs["output_scores"] = True
    generation_kwargs["return_dict_in_generate"] = True

    # Do inference
    with torch.no_grad():
        model_inputs = self.processor(images, return_tensors="pt").pixel_values
        model_outputs = self.model.generate(model_inputs.to(self.model.device), **generation_kwargs)

        texts = self.processor.batch_decode(model_outputs.sequences, skip_special_tokens=True)
        scores = self._compute_sequence_scores(model_outputs)

    # Assemble and return a list of Result objects from the prediction outputs.
    # `texts` and `scores` are flattened lists so we need to iterate over them in steps.
    # This is done to ensure that the list of results correspond 1-to-1 with the list of images.
    results = []
    metadata = self.metadata | {"generation_kwargs": generation_kwargs}
    step = generation_kwargs["num_return_sequences"]
    for i in range(0, len(texts), step):
        texts_chunk = texts[i : i + step]
        scores_chunk = scores[i : i + step]
        result = Result.text_recognition_result(metadata, texts_chunk, scores_chunk)
        results.append(result)
    return results

def _predict(self, images: list[NumpyImage], **kwargs) -> list[Result]:
    """
    Satrn-specific prediction method

    This method is used by `predict()` and should typically not be
    called directly.

    Arguments:
        images: Input images
        kwargs: Additional keyword arguments that are forwarded to
            `mmocr.apis.TextRecInferencer.__call__()`.
    """
    outputs = self.model(
        images,
        batch_size=len(images),
        return_datasamples=False,
        progress_bar=False,
        **kwargs,
    )
    results = []
    for prediction in outputs["predictions"]:
        texts = prediction["text"]
        scores = prediction["scores"]
        results.append(Result.text_recognition_result(self.metadata, texts, scores))
    return results

def _predict(self, images: list[np.ndarray], **decode_kwargs) -> list[Result]:
    """
    PyLaia-specific prediction method: runs text recognition.

    Args:
        images (list[np.ndarray]):
            List of images as NumPy arrays (e.g., shape [H, W, C]).
        batch_size (int, optional):
            Batch size for decoding. Defaults to 1.
        reading_order (str, optional):
            Reading order for text recognition. Defaults to "LTR".
        num_workers (int, optional):
            Number of workers for parallel processing. Defaults to `multiprocessing.cpu_count()`.

    Returns:
        list[Result]:
            A list of Result objects containing recognized text and
            optionally confidence scores.
    """

    temperature = decode_kwargs.get("temperature", 1.0)
    batch_size = decode_kwargs.get("batch_size", 1)
    reading_order = decode_kwargs.get("reading_order", "LTR")
    num_workers = decode_kwargs.get("num_workers", multiprocessing.cpu_count())

    common_args = CommonArgs(
        checkpoint="weights.ckpt",
        train_path=str(self.model_dir),
        experiment_dirname="",
    )

    data_args = DataArgs(
        batch_size=batch_size, color_mode="L", reading_order=reading_order, num_workers=num_workers
    )

    gpus_flag = 1 if self.device.type == "cuda" else 0
    trainer_args = TrainerArgs(gpus=gpus_flag)

    decode_args = DecodeArgs(
        include_img_ids=True,
        join_string="",
        convert_spaces=True,
        print_line_confidence_scores=True,
        print_word_confidence_scores=False,
        temperature=temperature,
        use_language_model=self.use_language_model,
        **self.language_model_params.model_dump(),
    )

    # Note: PyLaia's 'decode' function expects disk-based file paths rather than in-memory data.
    # Because it is tightly integrated as a CLI tool, we must create temporary image files
    # and pass their paths to the PyLaia decoder. Otherwise, PyLaia cannot process these images.
    tmp_images_dir = Path(mkdtemp())
    logger.debug(f"Created temp folder for images: {tmp_images_dir}")

    image_ids = [str(uuid4()) for _ in images]

    for img_id, np_img in zip(image_ids, images):
        padded_img = _ensure_min_height(np_img, 128)  # Just to fix the min pixel height (defaults to 128)
        cv2.imwrite(str(tmp_images_dir / f"{img_id}.jpg"), padded_img)

    with NamedTemporaryFile() as pred_stdout, NamedTemporaryFile() as img_list:
        Path(img_list.name).write_text("\n".join(image_ids))

        with redirect_stdout(open(pred_stdout.name, mode="w")):
            decode(
                syms=str(self.model_dir / "syms.txt"),
                img_list=img_list.name,
                img_dirs=[str(tmp_images_dir)],
                common=common_args,
                data=data_args,
                trainer=trainer_args,
                decode=decode_args,
                num_workers=num_workers,
            )
            sys.stdout.flush()

        decode_output_lines = Path(pred_stdout.name).read_text().strip().splitlines()

    results = []
    metadata = self.metadata | {"decode_kwargs": decode_kwargs}

    for line in decode_output_lines:
        match = self.LINE_PREDICTION.match(line)
        if not match:
            logger.warning("Could not parse line: %s", line)
            continue
        _, score_str, text = match.groups()  # _ = image_id

        try:
            score_val = float(score_str)
        except ValueError:
            score_val = 0.0

        result = Result.text_recognition_result(metadata, [text], [score_val])
        results.append(result)

    logger.debug(f"PyLaia recognized {len(results)} lines of text.")

    return results

def _predict(
    self,
    images: list[NumpyImage],
    nms_downscale: float = 1.0,
    nms_threshold: float = 0.4,
    nms_sigma: float = 2.0,
    **kwargs,
) -> list[Result]:
    """
    RTMDet-specific prediction method

    This method is used by `predict()` and should typically not be
    called directly.

    Arguments:
        images: List of input images
        nms_downscale: If < 1, all masks will be downscaled by this factor
            before applying NMS. This leads to faster NMS at the expense of
            accuracy.
        nms_threshold: Score threshold for segments to keep after NMS.
        nms_sigma: NMS parameter that affects the score calculation.
        **kwargs: Additional arguments that are passed to DetInferencer.__call__.
    """
    batch_size = max(1, len(images))
    outputs = self.model(
        images,
        batch_size=batch_size,
        draw_pred=False,
        return_datasample=True,
        **kwargs,
    )
    results = []
    for image, output in zip(images, outputs["predictions"]):
        results.append(self._create_segmentation_result(image, output, nms_downscale, nms_threshold, nms_sigma))
    return results

def _predict(
    self, images: list[np.ndarray], use_polygons: bool = True, polygon_approx_level: float = 0.005, **kwargs
) -> list[Result]:
    """
    Run inference.

    Arguments:
        images: Input images
        use_polygons: Wheter to include output polygons (if available), default True.
        polygon_approx_level: A parameter which controls the maximum distance between the original polygon
            and the approximated low-resolution polygon, as a fraction of the original polygon arc length.
            Example: With `polygon_approx_level=0.005` and a generated polygon with arc length 100, the
            approximated polygon will not differ more than 0.5 units from the original.
        **kwargs: Keyword arguments forwarded to the inner YOLO model instance.
    """
    outputs = self.model(images, stream=True, verbose=False, **kwargs)

    results = []
    for image, output in zip(images, outputs):
        polygons = bboxes = scores = class_labels = None
        if output.boxes is not None:
            bboxes = output.boxes.xyxy.int().tolist()
            scores = output.boxes.conf.tolist()
            class_labels = [output.names[label] for label in output.boxes.cls.tolist()]

        if use_polygons:
            if output.masks is None:
                logger.warning("`use_polygons` was set to True but the model did not return any polygons.")
            else:
                polygons = _simplify_polygons(output.masks.xy, polygon_approx_level)

        result = Result.segmentation_result(
            image.shape[:2],
            bboxes=bboxes,
            polygons=polygons,
            scores=scores,
            labels=class_labels,
            metadata=self.metadata,
        )
        results.append(result)
    return results

def _predict(
    self,
    images: list[np.ndarray],
    return_format: Literal["argmax", "softmax"] = "softmax",
) -> list[Result]:
    """Perform inference on `images`

    Arguments:
        images: List of input images.
        return_format: Decides the format of the output. Options are:
            - "softmax": returns the confidence scores for each class
                label and image. Default.
            - "argmax": returns the most probable class label for each
                image.
    """
    inputs = self.processor(images, return_tensors="pt").pixel_values

    with torch.no_grad():
        batch_logits = self.model(inputs.to(self.model.device)).logits

    return [
        Result(
            metadata=self.metadata,
            data=[{"classification": self._get_labels(logits, return_format)}],
        )
        for logits in batch_logits
    ]