import logging import os import threading from collections.abc import Iterable from typing import Dict, List, Set, Union import numpy as np import torch from PIL import Image from torch import Tensor from transformers import AutoModelForObjectDetection, RTDetrImageProcessor from docling_ibm_models.layoutmodel.labels import LayoutLabels _log = logging.getLogger(__name__) # Global lock for model initialization to prevent threading issues _model_init_lock = threading.Lock() class LayoutPredictor: """ Document layout prediction using safe tensors """ def __init__( self, artifact_path: str, device: str = "cpu", num_threads: int = 4, base_threshold: float = 0.3, blacklist_classes: Set[str] = set(), ): """ Provide the artifact path that contains the LayoutModel file Parameters ---------- artifact_path: Path for the model torch file. device: (Optional) device to run the inference. num_threads: (Optional) Number of threads to run the inference if device = 'cpu' Raises ------ FileNotFoundError when the model's torch file is missing """ # Blacklisted classes self._black_classes = blacklist_classes # set(["Form", "Key-Value Region"]) # Canonical classes self._labels = LayoutLabels() # Set basic params self._threshold = base_threshold # Score threshold # Set number of threads for CPU self._device = torch.device(device) self._num_threads = num_threads if device == "cpu": torch.set_num_threads(self._num_threads) # Load model file and configurations self._processor_config = os.path.join(artifact_path, "preprocessor_config.json") self._model_config = os.path.join(artifact_path, "config.json") self._st_fn = os.path.join(artifact_path, "model.safetensors") if not os.path.isfile(self._st_fn): raise FileNotFoundError("Missing safe tensors file: {}".format(self._st_fn)) if not os.path.isfile(self._processor_config): raise FileNotFoundError( f"Missing processor config file: {self._processor_config}" ) if not os.path.isfile(self._model_config): raise FileNotFoundError(f"Missing model config file: {self._model_config}") # Load model and move to device self._image_processor = RTDetrImageProcessor.from_json_file( self._processor_config ) # Use lock to prevent threading issues during model initialization with _model_init_lock: self._model = AutoModelForObjectDetection.from_pretrained( artifact_path, config=self._model_config, device_map=self._device ) self._model.eval() # Set classes map self._model_name = type(self._model).__name__ if self._model_name == "RTDetrForObjectDetection": self._classes_map = self._labels.shifted_canonical_categories() self._label_offset = 1 else: self._classes_map = self._labels.canonical_categories() self._label_offset = 0 _log.debug("LayoutPredictor settings: {}".format(self.info())) def info(self) -> dict: """ Get information about the configuration of LayoutPredictor """ info = { "model_name": self._model_name, "safe_tensors_file": self._st_fn, "device": self._device.type, "num_threads": self._num_threads, "image_size": self._image_processor.size, "threshold": self._threshold, } return info @torch.inference_mode() def predict(self, orig_img: Union[Image.Image, np.ndarray]) -> Iterable[dict]: """ Predict bounding boxes for a given image. The origin (0, 0) is the top-left corner and the predicted bbox coords are provided as: [left, top, right, bottom] Parameter --------- origin_img: Image to be predicted as a PIL Image object or numpy array. Yield ----- Bounding box as a dict with the keys: "label", "confidence", "l", "t", "r", "b" Raises ------ TypeError when the input image is not supported """ # Convert image format if isinstance(orig_img, Image.Image): page_img = orig_img.convert("RGB") elif isinstance(orig_img, np.ndarray): page_img = Image.fromarray(orig_img).convert("RGB") else: raise TypeError("Not supported input image format") target_sizes = torch.tensor([page_img.size[::-1]]) inputs = self._image_processor(images=[page_img], return_tensors="pt").to( self._device ) outputs = self._model(**inputs) results: List[Dict[str, Tensor]] = ( self._image_processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=self._threshold, ) ) w, h = page_img.size result = results[0] for score, label_id, box in zip( result["scores"], result["labels"], result["boxes"] ): score = float(score.item()) label_id = int(label_id.item()) + self._label_offset label_str = self._classes_map[label_id] # Filter out blacklisted classes if label_str in self._black_classes: continue bbox_float = [float(b.item()) for b in box] l = min(w, max(0, bbox_float[0])) t = min(h, max(0, bbox_float[1])) r = min(w, max(0, bbox_float[2])) b = min(h, max(0, bbox_float[3])) yield { "l": l, "t": t, "r": r, "b": b, "label": label_str, "confidence": score, } @torch.inference_mode() def predict_batch( self, images: List[Union[Image.Image, np.ndarray]] ) -> List[List[dict]]: """ Batch prediction for multiple images - more efficient than calling predict() multiple times. Parameters ---------- images : List[Union[Image.Image, np.ndarray]] List of images to process in a single batch Returns ------- List[List[dict]] List of prediction lists, one per input image. Each prediction dict contains: "label", "confidence", "l", "t", "r", "b" """ if not images: return [] # Convert all images to RGB PIL format pil_images = [] for img in images: if isinstance(img, Image.Image): pil_images.append(img.convert("RGB")) elif isinstance(img, np.ndarray): pil_images.append(Image.fromarray(img).convert("RGB")) else: raise TypeError("Not supported input image format") # Get target sizes for all images target_sizes = torch.tensor([img.size[::-1] for img in pil_images]) # Process all images in a single batch inputs = self._image_processor(images=pil_images, return_tensors="pt").to( self._device ) outputs = self._model(**inputs) # Post-process all results at once results_list: List[Dict[str, Tensor]] = ( self._image_processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=self._threshold, ) ) # Convert results to standard format for each image all_predictions = [] for img, results in zip(pil_images, results_list): w, h = img.size predictions = [] for score, label_id, box in zip( results["scores"], results["labels"], results["boxes"] ): score = float(score.item()) label_id = int(label_id.item()) + self._label_offset label_str = self._classes_map[label_id] # Filter out blacklisted classes if label_str in self._black_classes: continue bbox_float = [float(b.item()) for b in box] l = min(w, max(0, bbox_float[0])) t = min(h, max(0, bbox_float[1])) r = min(w, max(0, bbox_float[2])) b = min(h, max(0, bbox_float[3])) predictions.append( { "l": l, "t": t, "r": r, "b": b, "label": label_str, "confidence": score, } ) all_predictions.append(predictions) return all_predictions