Smarttext

class SmartTextBackbone(StrEnum):
    """Backbone names supported by the original SmartText scorer."""

    shufflenetv2 = "shufflenetv2"
    mobilenetv2 = "mobilenetv2"
    vgg16 = "vgg16"
    resnet50 = "resnet50"

class SmartTextConfig(PretrainedConfig):
    """Configuration for SmartText scorer, saliency model, and pipeline.

    Args:
        id2label: Public label mapping. Defaults to one ``text`` label.
        scorer_scale: Original scorer scale mode.
        scorer_backbone: Original scorer backbone.
        align_size: RoI/RoD pooled spatial size.
        reduction_dim: Reduced feature channels before scoring.
        downsample: Scorer feature-map downsampling factor.
        model_type_name: Original ``RoE``/``RoD`` region mode.
        image_size: Short-side normalization target for scorer preprocessing.
        ratio_list: Per-line font-size ratios.
        text_spacing: Pixel spacing between prompt lines.
        exp_prop: Original expanded-region coefficient.
        grid_num: Candidate search grid count.
        saliency_coef: Saliency suppression coefficient.
        max_text_area_coef: Maximum candidate area divisor.
        min_text_area_coef: Minimum candidate area divisor.
        min_font_size: Minimum candidate font size.
        max_font_size: Maximum candidate font size.
        font_inc_unit: Candidate font-size step.
        candi_res: Number of selected candidates.
        contrast_threshold: Foreground/background contrast threshold.
        mos_mean: MOS score mean used by the original demo.
        mos_std: MOS score standard deviation.
        rgb_mean: RGB normalization mean for scorer inputs.
        rgb_std: RGB normalization std for scorer inputs.
        scorer_subfolder: Pipeline scorer subfolder.
        saliency_subfolder: Pipeline saliency-model subfolder.
        processor_subfolder: Pipeline processor subfolder.
        original_options: Raw reference option values preserved for audit.
        conversion_report: Conversion metadata persisted in configs.
        kwargs: Extra ``PretrainedConfig`` fields.

    Examples:
        >>> config = SmartTextConfig()
        >>> config.id2label
        {0: 'text'}
    """

    model_type = "smarttext"

    def __init__(
        self,
        *,
        id2label: Mapping[int | str, str] | None = None,
        scorer_scale: str = "multi",
        scorer_backbone: SmartTextBackbone | str = SmartTextBackbone.shufflenetv2,
        align_size: int = 9,
        reduction_dim: int = 8,
        downsample: int = 4,
        model_type_name: SmartTextRegionMode | str = SmartTextRegionMode.RoE,
        image_size: int = 256,
        ratio_list: Sequence[float] = (1.0, 0.8),
        text_spacing: int = 20,
        exp_prop: int = 6,
        grid_num: int = 120,
        saliency_coef: float = 2.6,
        max_text_area_coef: float = 17.0,
        min_text_area_coef: float = 7.0,
        min_font_size: int = 10,
        max_font_size: int = 500,
        font_inc_unit: int = 5,
        candi_res: int = 3,
        contrast_threshold: float = 5.0,
        mos_mean: float = 2.95,
        mos_std: float = 0.8,
        rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
        rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
        scorer_subfolder: str = "scorer",
        saliency_subfolder: str = "saliency_model",
        processor_subfolder: str = "processor",
        original_options: Mapping[str, SmartTextMetadataValue] | None = None,
        conversion_report: Mapping[str, SmartTextMetadataValue | list[str]]
        | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize SmartText configuration."""
        raw_id2label = id2label or DEFAULT_ID2LABEL
        normalized_id2label = {int(key): value for key, value in raw_id2label.items()}
        super().__init__(id2label=normalized_id2label, **kwargs)  # ty: ignore[invalid-argument-type]

        self.id2label = normalized_id2label
        self.label2id = {value: key for key, value in self.id2label.items()}

        self.scorer_scale = scorer_scale
        self.scorer_backbone = SmartTextBackbone(scorer_backbone).value
        self.align_size = int(align_size)
        self.reduction_dim = int(reduction_dim)
        self.downsample = int(downsample)

        self.model_type_name = SmartTextRegionMode(model_type_name).value
        self.image_size = int(image_size)
        self.ratio_list = tuple(float(value) for value in ratio_list)
        self.text_spacing = int(text_spacing)
        self.exp_prop = int(exp_prop)

        self.grid_num = int(grid_num)
        self.saliency_coef = float(saliency_coef)
        self.max_text_area_coef = float(max_text_area_coef)
        self.min_text_area_coef = float(min_text_area_coef)
        self.min_font_size = int(min_font_size)
        self.max_font_size = int(max_font_size)
        self.font_inc_unit = int(font_inc_unit)
        self.candi_res = int(candi_res)
        self.contrast_threshold = float(contrast_threshold)
        self.mos_mean = float(mos_mean)
        self.mos_std = float(mos_std)
        self.rgb_mean = tuple(float(value) for value in rgb_mean)
        self.rgb_std = tuple(float(value) for value in rgb_std)

        self.scorer_subfolder = scorer_subfolder
        self.saliency_subfolder = saliency_subfolder
        self.processor_subfolder = processor_subfolder

        self.original_options = dict(original_options or {})
        self.conversion_report = dict(conversion_report or {})

    @property
    def num_labels(self) -> int:
        """Return the number of public semantic labels."""
        return len(cast(dict[int, str], self.id2label))

    @property
    def uses_expanded_region(self) -> bool:
        """Return whether the original ``RoE`` expanded-region mode is active."""
        return self.model_type_name == SmartTextRegionMode.RoE.value

    @property
    def scorer_input_channels(self) -> int:
        """Return the RGB scorer input channel count."""
        return 3

class SmartTextRegionMode(StrEnum):
    """Supported SmartText region scoring modes."""

    RoD = "RoD"
    RoE = "RoE"

class SmartTextImageProcessor(BaseImageProcessor):
    """Prepare SmartText scorer and BASNet image tensors.

    Args:
        image_size: Scorer short-side target.
        rgb_mean: Scorer RGB normalization mean.
        rgb_std: Scorer RGB normalization standard deviation.

    Examples:
        >>> processor = SmartTextImageProcessor()
        >>> batch = processor.preprocess(Image.new("RGB", (32, 32)))
        >>> tuple(batch["pixel_values"].shape[:2])
        (1, 3)
    """

    model_input_names = ["pixel_values", "basnet_pixel_values"]

    def __init__(
        self,
        image_size: int = 256,
        rgb_mean: Sequence[float] = (0.485, 0.456, 0.406),
        rgb_std: Sequence[float] = (0.229, 0.224, 0.225),
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize image processor settings."""
        super().__init__(**kwargs)
        self.image_size = int(image_size)
        self.rgb_mean = tuple(float(value) for value in rgb_mean)
        self.rgb_std = tuple(float(value) for value in rgb_std)

    @classmethod
    def from_config(cls, config: SmartTextConfig) -> "SmartTextImageProcessor":
        """Build an image processor from SmartText configuration."""
        return cls(
            image_size=config.image_size,
            rgb_mean=config.rgb_mean,
            rgb_std=config.rgb_std,
        )

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput],
        *,
        return_tensors: Literal["pt"] = "pt",
        target_min_side: int | None = None,
        rgb_mean: Sequence[float] | None = None,
        rgb_std: Sequence[float] | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> BatchFeature:
        """Preprocess images for the SmartText scorer.

        Args:
            images: RGB image or image batch.
            return_tensors: Tensor framework. Only ``pt`` is supported.
            target_min_side: Optional short-side target override.
            rgb_mean: Optional RGB mean override.
            rgb_std: Optional RGB std override.
            kwargs: Ignored compatibility kwargs.

        Returns:
            Batch feature with ``pixel_values`` and ``image_sizes``.

        Raises:
            ValueError: If ``return_tensors`` is not ``pt``.
        """
        del kwargs
        if return_tensors != "pt":
            raise ValueError(
                "SmartTextImageProcessor only supports return_tensors='pt'"
            )

        mean = np.asarray(rgb_mean or self.rgb_mean, dtype=np.float32)
        std = np.asarray(rgb_std or self.rgb_std, dtype=np.float32)
        tensors = []
        sizes = []
        for image in _ensure_pil_batch(images):
            width, height = image.size
            sizes.append((height, width))
            scale = (target_min_side or self.image_size) / min(height, width)
            resized_h = max(32, int(round(height * scale / 32.0) * 32))
            resized_w = max(32, int(round(width * scale / 32.0) * 32))
            resized = image.convert("RGB").resize(
                (resized_w, resized_h), Image.Resampling.BILINEAR
            )
            array = np.asarray(resized, dtype=np.float32) / 256.0
            tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))
        return BatchFeature(
            {
                "pixel_values": torch.stack(tensors).float(),
                "image_sizes": torch.tensor(sizes, dtype=torch.long),
            }
        )

    def preprocess_basnet(
        self,
        images: ImageInput | Sequence[ImageInput],
        *,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchFeature:
        """Preprocess images for BASNet saliency prediction.

        Args:
            images: RGB image or image batch.
            return_tensors: Tensor framework. Only ``pt`` is supported.

        Returns:
            Batch feature with ``basnet_pixel_values`` shaped ``(B, 3, 256, 256)``.
        """
        basnet = BASNetImageProcessor(
            input_size=256,
            rgb_mean=self.rgb_mean,
            rgb_std=self.rgb_std,
        ).preprocess(images, return_tensors=return_tensors)
        return BatchFeature(
            {
                "basnet_pixel_values": basnet["pixel_values"],
                "image_sizes": basnet["image_sizes"],
            }
        )

class SmartTextBASNet(BASNetModel):
    """SmartText saliency component backed by the shared BASNet model."""

    config_class = SmartTextConfig

    def __init__(self, config: SmartTextConfig) -> None:
        """Initialize the SmartText saliency component."""
        super().__init__(cast(BASNetConfig, config))

class SmartTextScorer(PreTrainedModel):
    """Reference-compatible SMT candidate scorer.

    The module names match the original ``smtModel.py`` for
    ``build_smt_model(scale="multi", alignsize=9, reddim=8,
    model="shufflenetv2", downsample=4)`` so ``SMT.pth`` loads directly.

    Args:
        config: SmartText configuration.

    Examples:
        >>> config = SmartTextConfig()
        >>> model = SmartTextScorer(config)
        >>> "Feat_ext.feature3.0.0.weight" in model.state_dict()
        True
    """

    config_class = SmartTextConfig
    main_input_name = "pixel_values"
    _tied_weights_keys: list[str] = []

    def __init__(self, config: SmartTextConfig) -> None:
        """Initialize scorer scorer modules."""
        super().__init__(config)
        self.all_tied_weights_keys: dict[str, str] = {}
        if config.scorer_scale != "multi" or config.scorer_backbone != "shufflenetv2":
            raise ValueError(
                "Only the released multi-scale ShuffleNetV2 SMT scorer is supported"
            )

        self.Feat_ext = _ShuffleNetV2Base()
        self.DimRed = nn.Conv2d(812, config.reduction_dim, kernel_size=1, padding=0)
        self.downsample2 = nn.UpsamplingBilinear2d(scale_factor=1.0 / 2.0)
        self.upsample2 = nn.UpsamplingBilinear2d(scale_factor=2.0)
        spatial_scale = 1.0 / 2**config.downsample
        self.RoIAlign = SmartTextRoIAlignAvg(
            config.align_size + 1, config.align_size + 1, spatial_scale
        )
        self.RoDAlign = SmartTextRoDAlignAvg(
            config.align_size + 1, config.align_size + 1, spatial_scale
        )
        self.FC_layers = _fc_layers(config.reduction_dim * 2, config.align_size)

    def forward(
        self,
        pixel_values: Float[torch.Tensor, "batch channels height width"],
        boxes: Float[torch.Tensor, "candidates 5"],
        return_dict: bool | None = None,
    ) -> SmartTextScorerOutput | tuple[Float[torch.Tensor, "candidates"]]:
        """Score candidate text regions.

        Args:
            pixel_values: RGB scorer tensor shaped ``(B, 3, H, W)``.
            boxes: RoI rows shaped ``(N, 5)`` with batch index in column zero.
            return_dict: Whether to return a ``ModelOutput``.

        Returns:
            Candidate scores as a ``SmartTextScorerOutput`` or tuple.
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        f3, f4, f5 = self.Feat_ext(pixel_values)
        cat_feat = torch.cat((self.downsample2(f3), f4, 0.5 * self.upsample2(f5)), 1)
        red_feat = self.DimRed(cat_feat)
        roi_feat = self.RoIAlign(red_feat, boxes)
        rod_feat = self.RoDAlign(red_feat, boxes)
        prediction = self.FC_layers(torch.cat((roi_feat, rod_feat), 1)).flatten()
        if not return_dict:
            return (prediction,)
        return SmartTextScorerOutput(scores=prediction)

@dataclass
class SmartTextScorerOutput(ModelOutput):
    """Output of ``SmartTextScorer.forward``."""

    scores: Float[torch.Tensor, "candidates"]

class SmartTextPipeline(LayoutGenerationPipeline):
    """Transformers-side SmartText pipeline.

    Args:
        scorer: Candidate scoring model.
        saliency_model: BASNet saliency model.
        processor: Input/output processor.
        config: Optional root pipeline config.
        device: Optional runtime device.

    Examples:
        >>> config = SmartTextConfig(align_size=3, reduction_dim=4, grid_num=16, max_font_size=20)
        >>> pipe = SmartTextPipeline(SmartTextScorer(config), SmartTextBASNet(config), config=config)
        >>> pipe.config.model_type
        'smarttext'
    """

    config_class: ClassVar[type[PretrainedConfig]] = SmartTextConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = {
        "scorer": PipelineComponentSpec(
            attribute_name="scorer",
            loader=_load_scorer_component,
            marker_file="config.json",
            config_subfolder_attribute="scorer_subfolder",
        ),
        "saliency_model": PipelineComponentSpec(
            attribute_name="saliency_model",
            loader=_load_saliency_component,
            marker_file="config.json",
            config_subfolder_attribute="saliency_subfolder",
        ),
        "processor": PipelineComponentSpec(
            attribute_name="processor",
            loader=_load_processor_component,
            marker_file="processor_config.json",
            save_with_is_main_process=False,
            config_subfolder_attribute="processor_subfolder",
        ),
    }

    config: SmartTextConfig
    scorer: SmartTextScorer
    saliency_model: SmartTextBASNet
    processor: SmartTextProcessor

    def __init__(
        self,
        scorer: SmartTextScorer,
        saliency_model: SmartTextBASNet,
        processor: SmartTextProcessor | None = None,
        config: SmartTextConfig | None = None,
        device: str | torch.device | None = None,
    ) -> None:
        """Initialize SmartText pipeline."""
        super().__init__(config or scorer.config)
        self.config = config or scorer.config
        self.scorer = scorer
        self.saliency_model = saliency_model
        self.processor = processor or SmartTextProcessor(config=self.config)
        self.scorer.eval()
        self.saliency_model.eval()
        if device is not None:
            self.to(device)

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> "SmartTextPipeline":
        """Build a SmartText pipeline from loaded components."""
        return cls(
            config=cast(SmartTextConfig, config),
            scorer=cast(SmartTextScorer, components["scorer"]),
            saliency_model=cast(SmartTextBASNet, components["saliency_model"]),
            processor=cast(SmartTextProcessor, components["processor"]),
        )

    @torch.no_grad()
    def __call__(  # ty: ignore[invalid-method-override]
        self,
        images: ImageInput
        | Sequence[ImageInput]
        | Float[torch.Tensor, "batch channels height width"]
        | None = None,
        *,
        content: Mapping[
            str,
            ImageInput
            | Sequence[ImageInput]
            | str
            | Sequence[str]
            | Float[torch.Tensor, "batch height width"]
            | Sequence[CandidateBoxRow]
            | Sequence[Sequence[CandidateBoxRow]],
        ]
        | None = None,
        prompt: str | Sequence[str] | None = None,
        text: str | Sequence[str] | None = None,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.content_image,
        labels: Int[torch.Tensor, "batch elements"]
        | Sequence[Sequence[int]]
        | Sequence[int]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Sequence[Sequence[Sequence[float]]]
        | Sequence[Sequence[float]]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Sequence[Sequence[bool]]
        | Sequence[bool]
        | None = None,
        num_elements: int | Sequence[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        output_type: OutputType | Literal["dataclass", "dict"] = OutputType.dataclass,
        return_intermediates: bool = False,
        font: str | Path | ImageFont.FreeTypeFont | ImageFont.ImageFont | None = None,
        ratio_list: Sequence[float] | None = None,
        text_spacing: int | None = None,
        candi_res: int | None = None,
        saliency: ImageInput
        | Sequence[ImageInput]
        | Float[torch.Tensor, "batch height width"]
        | None = None,
        candidate_boxes: Sequence[CandidateBoxRow]
        | Sequence[Sequence[CandidateBoxRow]]
        | None = None,
        return_text_lines: bool = False,
        score_normalization: Literal["mos", "raw"] = "mos",
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | dict[
                str,
                Shaped[torch.Tensor, "..."]
                | str
                | list[int]
                | list[SmartTextCandidate]
                | None,
            ]
            | None,
        ]
    ):
        """Generate text placement boxes for content images.

        Args:
            images: RGB image or image batch.
            content: Optional content carrier.
            prompt: Prompt text payload.
            text: Alias for prompt.
            batch_size: Expected batch size for validation.
            seed: Optional seed used only when ``generator`` is absent.
            generator: Explicit torch generator; wins over ``seed``.
            condition_type: Must normalize to ``content_image``.
            labels: Unsupported v1 compatibility argument.
            bbox: Unsupported v1 compatibility argument.
            mask: Unsupported v1 compatibility argument.
            num_elements: Unsupported v1 compatibility argument.
            box_format: Public v1 compatibility argument.
            normalized: Public v1 compatibility argument.
            canvas_size: Optional source canvas size override.
            num_inference_steps: Unused v1 compatibility argument.
            output_type: ``dataclass`` or ``dict``.
            return_intermediates: Whether to include intermediate payloads.
            font: TrueType font path or PIL font object.
            ratio_list: Optional per-line font ratios.
            text_spacing: Optional text-spacing override.
            candi_res: Optional top-k override.
            saliency: Optional saliency map bypassing BASNet.
            candidate_boxes: Optional reference-style candidates.
            return_text_lines: Return per-line boxes for top candidate.
            score_normalization: ``mos`` or ``raw``.

        Returns:
            Shared layout output.
        """
        del (
            labels,
            bbox,
            mask,
            num_elements,
            box_format,
            normalized,
            canvas_size,
            num_inference_steps,
        )
        normalize_condition_type(condition_type)
        self.prepare_generator(generator=generator, seed=seed, device=self.device)
        if font is None:
            font = ImageFont.load_default()
        effective_config = self.config
        if text_spacing is not None:
            effective_config = copy.copy(self.config)
            effective_config.text_spacing = int(text_spacing)
        encoded = self.processor(
            images,
            content=content,
            prompt=prompt,
            text=text,
            saliency=saliency,
            candidate_boxes=candidate_boxes,
            font=font,
        )
        if len(encoded["images"]) != 1 or batch_size != 1:
            raise ValueError("SmartText currently decodes one image at a time")

        image = encoded["images"][0]
        prompt_text = encoded["prompts"][0]
        resolved_saliency = encoded["saliency"]
        if resolved_saliency is None:
            basnet_values = encoded["basnet_pixel_values"].to(self._runtime_device())
            saliency_out = self.saliency_model(basnet_values)
            resolved_saliency = saliency_out.saliency[0]
        candidates = encoded["candidate_boxes"]
        if candidates is None:
            candidates = generate_candidates(
                image,
                cast(torch.Tensor, resolved_saliency),
                prompt=prompt_text,
                font=font,
                config=effective_config,
                ratio_list=ratio_list,
            )
        pixel_values, boxes, candidates = prepare_scorer_batch(
            image,
            candidates,
            config=effective_config,
        )
        scores = self.scorer(
            pixel_values.to(self._runtime_device()),
            boxes.to(self._runtime_device()),
        ).scores
        top_k = candi_res or self.config.candi_res
        raw_scores = scores.detach().cpu().float().flatten()
        if candidates:
            selected_color_index = max(
                range(len(candidates)), key=lambda index: float(raw_scores[index])
            )
            text_color = choose_text_color(
                image,
                candidates[selected_color_index].bbox_ltrb_px,
                contrast_threshold=self.config.contrast_threshold,
            )
        else:
            text_color = None
        intermediates = None
        if return_intermediates:
            intermediates = {
                "saliency": resolved_saliency,
                "raw_scorer_boxes": boxes.detach().cpu(),
                "prompt": prompt_text,
            }
        return self.processor.decode(
            candidates=candidates,
            scores=scores,
            image_size=image.size,
            output_type=cast(Literal["dataclass", "dict"], str(output_type)),
            return_text_lines=return_text_lines,
            top_k=top_k,
            score_normalization=score_normalization,
            text_color=text_color,
            intermediates=intermediates,
        )

    def _runtime_device(self) -> torch.device:
        return self.device or next(self.scorer.parameters()).device