processing_astroclip.py

# Copyright 2026 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.

"""Processor class for AstroCLIP."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Optional

from transformers.dynamic_module_utils import custom_object_save
from transformers.processing_utils import ProcessorMixin
from transformers.utils import PROCESSOR_NAME, cached_file

from .feature_extraction_astroclip import AstroClipSpectrumFeatureExtractor
from .image_processing_astroclip import AstroClipImageProcessor


class AstroClipProcessor(ProcessorMixin):
    """Compose AstroCLIP image and spectrum preprocessing."""

    attributes = ["image_processor", "spectrum_feature_extractor"]
    image_processor_class = "AstroClipImageProcessor"
    spectrum_feature_extractor_class = "AstroClipSpectrumFeatureExtractor"

    def __init__(
        self,
        image_processor: Optional[AstroClipImageProcessor] = None,
        spectrum_feature_extractor: Optional[AstroClipSpectrumFeatureExtractor] = None,
        spectrum_processor: Optional[AstroClipSpectrumFeatureExtractor] = None,
        **kwargs,
    ):
        if kwargs:
            unexpected = ", ".join(sorted(kwargs))
            raise TypeError(f"unexpected AstroCLIP processor arguments: {unexpected}")
        if spectrum_processor is not None:
            if spectrum_feature_extractor is not None:
                raise ValueError("pass only one of spectrum_feature_extractor or spectrum_processor")
            spectrum_feature_extractor = spectrum_processor
        if image_processor is None:
            image_processor = AstroClipImageProcessor()
        if spectrum_feature_extractor is None:
            spectrum_feature_extractor = AstroClipSpectrumFeatureExtractor()
        # Older Transformers releases validate processor components through
        # global AutoClass mappings, which are not reliable for remote code.
        self.image_processor = image_processor
        self.spectrum_feature_extractor = spectrum_feature_extractor

    @property
    def spectrum_processor(self) -> AstroClipSpectrumFeatureExtractor:
        return self.spectrum_feature_extractor

    def to_dict(self) -> dict[str, Any]:
        image_processor = self.image_processor
        spectrum_feature_extractor = self.spectrum_feature_extractor
        output = {
            "processor_class": self.__class__.__name__,
            "image_processor": {
                "crop_size": image_processor.crop_size,
                "bands": image_processor.bands,
                "band_indices": image_processor.band_indices,
                "m": image_processor.m,
                "q": image_processor.q,
            },
            "spectrum_feature_extractor": {
                "section_length": spectrum_feature_extractor.section_length,
                "overlap": spectrum_feature_extractor.overlap,
                "min_std": spectrum_feature_extractor.min_std,
            },
        }
        if self._auto_class is not None:
            output["auto_map"] = {self._auto_class: "processing_astroclip.AstroClipProcessor"}
        return output

    def save_pretrained(self, save_directory: str | Path, **kwargs):
        save_directory = Path(save_directory)
        save_directory.mkdir(parents=True, exist_ok=True)
        processor_dict = self.to_dict()
        if self._auto_class is not None:
            custom_object_save(self, save_directory, config=processor_dict)
        processor_path = save_directory / PROCESSOR_NAME
        processor_path.write_text(
            json.dumps(processor_dict, indent=2) + "\n",
            encoding="utf-8",
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        local_path = Path(pretrained_model_name_or_path) / PROCESSOR_NAME
        if local_path.exists():
            processor_file = local_path
        else:
            hub_kwargs = {
                key: kwargs[key]
                for key in (
                    "cache_dir",
                    "force_download",
                    "local_files_only",
                    "token",
                    "revision",
                    "subfolder",
                    "repo_type",
                    "user_agent",
                )
                if key in kwargs
            }
            processor_file = cached_file(
                pretrained_model_name_or_path,
                PROCESSOR_NAME,
                **hub_kwargs,
            )
        if processor_file is None:
            raise OSError(f"could not find {PROCESSOR_NAME} in {pretrained_model_name_or_path}")
        processor_dict = json.loads(Path(processor_file).read_text(encoding="utf-8"))
        return cls(
            image_processor=AstroClipImageProcessor(**processor_dict.get("image_processor", {})),
            spectrum_feature_extractor=AstroClipSpectrumFeatureExtractor(
                **processor_dict.get("spectrum_feature_extractor", {})
            ),
        )

    def __call__(
        self,
        images: Optional[Any] = None,
        spectra: Optional[Any] = None,
        return_tensors: Optional[str] = None,
        **kwargs,
    ):
        if images is None and spectra is None:
            raise ValueError("provide images, spectra or both")

        encoding = {}
        if images is not None:
            if self.image_processor is None:
                raise ValueError("an image_processor is required when images are provided")
            image_encoding = self.image_processor(
                images=images,
                return_tensors=return_tensors,
                **kwargs,
            )
            encoding.update(dict(image_encoding))
        if spectra is not None:
            spectrum_encoding = self.spectrum_feature_extractor(
                spectra,
                return_tensors=return_tensors,
            )
            encoding.update(dict(spectrum_encoding))
        return encoding


__all__ = ["AstroClipProcessor"]