Hugging Face compatibility

The student PIQA model must be HF-compatible:
- https://huggingface.co/docs/transformers/en/custom_models

---

Below is a complete implementation of a custom Hugging Face–compatible student model for panorama IQA. This implementation includes a custom configuration, image processor, and model class that wraps a Swin‑v2 backbone. The model supports two types of heads (“point” and “piven”) and registers itself with the Hugging Face auto classes so that it can be loaded via, for example, `AutoModel.from_pretrained`. You can adjust the default parameters (such as the backbone model name and head type) as needed.

```python
import torch
import torch.nn as nn
from transformers import PretrainedConfig, PreTrainedModel, AutoModel
from transformers import ImageProcessingMixin

#############################
# 1. Custom Configuration
#############################

class PanoramaStudentConfig(PretrainedConfig):
model_type = "panorama_student"

def __init__(
self,
backbone_model_name: str = "microsoft/swinv2-base-patch4-window8-256",
head_type: str = "point", # options: "point" or "piven"
**kwargs,
):
"""
Args:
backbone_model_name (str): The Hugging Face model id for the panorama backbone.
head_type (str): Type of head to use; "point" returns a single score, "piven" returns prediction intervals.
"""
super().__init__(**kwargs)
self.backbone_model_name = backbone_model_name
self.head_type = head_type
# You can add other hyperparameters as needed

#############################
# 2. Custom Image Processor
#############################

class PanoramaImageProcessor(ImageProcessingMixin):
model_input_names = ["pixel_values"]

def __init__(self, size: int = 256, mean=None, std=None):
"""
Simple image processor that resizes images, converts them to tensors,
and normalizes them. Adjust as necessary for your training regime.
"""
self.size = size
# Default normalization values (example for ImageNet-like statistics)
self.mean = mean if mean is not None else [0.485, 0.456, 0.406]
self.std = std if std is not None else [0.229, 0.224, 0.225]

def __call__(self, images, return_tensors="pt", **kwargs):
"""
Args:
images (list[PIL.Image.Image] or a single PIL.Image.Image): Input images.
return_tensors (str): "pt" for PyTorch tensors.
Returns:
A dict with the key "pixel_values".
"""
if not isinstance(images, (list, tuple)):
images = [images]

processed = [self.preprocess(image) for image in images]
pixel_values = torch.stack(processed)
if return_tensors == "pt":
return {"pixel_values": pixel_values}
else:
return {"pixel_values": pixel_values.cpu().numpy()}

def preprocess(self, image):
# Import here to avoid dependency issues if not used elsewhere
from PIL import Image
import torchvision.transforms as T

transform = T.Compose([
T.Resize((self.size, self.size)),
T.ToTensor(),
T.Normalize(mean=self.mean, std=self.std),
])
return transform(image)

#############################
# 3. Custom Student Model
#############################

class PanoramaStudentModel(PreTrainedModel):
config_class = PanoramaStudentConfig

def __init__(self, config: PanoramaStudentConfig):
"""
The student model uses only the panorama backbone. The overall feature map
is taken from the backbone and fed to a head that can be a point estimator
or a prediction intervals head (PIVEN).
"""
super().__init__(config)

# Load the panorama backbone (e.g., Swin-v2) from HF.
self.backbone = AutoModel.from_pretrained(config.backbone_model_name)
hidden_size = self.backbone.config.hidden_size

# Initialize the head depending on the chosen type.
if config.head_type == "point":
# For point estimation, output a single value.
self.head = nn.Sequential(
nn.Linear(hidden_size, 1),
nn.Sigmoid() # scales output to [0,1]; adjust if you prefer a different range.
)
elif config.head_type == "piven":
# For PIVEN, output three numbers:
# - lower bound, upper bound, and a weight (normalized with sigmoid)
self.head = nn.Linear(hidden_size, 3)
else:
raise ValueError(f"Unsupported head type: {config.head_type}")

# Initialize weights and apply any model-specific initialization if needed.
self.post_init()

def forward(self, pixel_values, **kwargs):
"""
Args:
pixel_values (torch.FloatTensor): Batch of preprocessed images, shape (batch_size, channels, height, width).
Returns:
dict: Contains the model output(s). For "point" head: {"score": score}.
For "piven" head: {"lower_bound": lb, "upper_bound": ub, "weight": weight, "point_estimate": point}.
"""
# Forward pass through the backbone.
backbone_outputs = self.backbone(pixel_values=pixel_values, **kwargs)
# Extract the panorama feature vector.
# (Assuming the backbone returns last_hidden_state of shape (batch, seq_len, hidden_size) and using the first token.)
features = backbone_outputs.last_hidden_state[:, 0, :]

if self.config.head_type == "point":
# Compute the quality score.
score = self.head(features)
return {"score": score}
else: # "piven" head
head_output = self.head(features)
# Split the outputs: first two entries for bounds, third for weight.
lower_bound = head_output[:, 0:1]
upper_bound = head_output[:, 1:2]
weight = torch.sigmoid(head_output[:, 2:3])
# Recover point estimate from the interval: point = lower_bound + weight * (upper_bound - lower_bound)
point_estimate = lower_bound + weight * (upper_bound - lower_bound)
return {
"lower_bound": lower_bound,
"upper_bound": upper_bound,
"weight": weight,
"point_estimate": point_estimate,
}

#############################
# 4. Registration for Auto Classes
#############################

# Note: Registration is needed only if you wish to load the model using the Auto classes.
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
from transformers.models.auto.modeling_auto import MODEL_MAPPING
from transformers.models.auto.processing_auto import PROCESSOR_MAPPING

# Register our custom config, model, and image processor.
CONFIG_MAPPING.register(PanoramaStudentConfig.model_type, PanoramaStudentConfig)
MODEL_MAPPING.register(PanoramaStudentConfig, PanoramaStudentModel)
PROCESSOR_MAPPING.register(PanoramaStudentConfig, PanoramaImageProcessor)

#############################
# Example: Saving and Loading
#############################

if __name__ == "__main__":
# Create a configuration for a point-estimate head.
config = PanoramaStudentConfig(
backbone_model_name="microsoft/swinv2-base-patch4-window8-256",
head_type="point",
)
# Instantiate the model.
model = PanoramaStudentModel(config)

# Create an image processor.
processor = PanoramaImageProcessor(size=256)

# Example: process an image (requires PIL)
# from PIL import Image
# image = Image.open("path/to/your/image.jpg")
# inputs = processor(image, return_tensors="pt")
# outputs = model(**inputs)
# print(outputs)

# Save the model and processor for later use.
model.save_pretrained("./panorama_student_model")
processor.save_pretrained("./panorama_student_model")

# Later, you can load them via the Auto classes:
# from transformers import AutoModel, AutoProcessor
# loaded_model = AutoModel.from_pretrained("./panorama_student_model")
# loaded_processor = AutoProcessor.from_pretrained("./panorama_student_model")
```

---

### Explanation

1. **Custom Configuration (`PanoramaStudentConfig`)**
This class inherits from `PretrainedConfig` and defines two key parameters: the backbone model (a Swin‑v2 model from HF) and the type of head ("point" or "piven").

2. **Custom Image Processor (`PanoramaImageProcessor`)**
Inherits from `ImageProcessingMixin` and provides a basic implementation to resize, convert, and normalize images. You may customize the transforms as needed.

3. **Custom Model (`PanoramaStudentModel`)**
Inherits from `PreTrainedModel` and uses the backbone to extract a panorama feature vector. Depending on the configuration, it passes the feature vector to either a point head (a linear layer with sigmoid) or a PIVEN head (a linear layer producing three outputs with a sigmoid applied to the weight).

4. **Registration with Auto Classes**
The config, model, and processor are registered with the HF auto mappings, so users can load your model using `AutoModel.from_pretrained` and `AutoProcessor.from_pretrained`.

This complete implementation makes the student model fully HF-compatible and ready for integration, training, or sharing within the Hugging Face ecosystem.