mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-03-12 12:25:32 +08:00
78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
"""Pydantic data models for feature compressor configuration."""
|
|
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
|
|
|
|
|
class PoolingType(str, Enum):
|
|
"""Enum for pooling types."""
|
|
|
|
ATTENTION = "attention"
|
|
|
|
|
|
class ModelConfig(BaseModel):
|
|
"""Configuration for the vision model and compression."""
|
|
|
|
model_config = ConfigDict(extra="ignore")
|
|
|
|
name: str = "facebook/dinov2-large"
|
|
compression_dim: int = Field(
|
|
default=256, gt=0, description="Output feature dimension"
|
|
)
|
|
pooling_type: PoolingType = PoolingType.ATTENTION
|
|
top_k_ratio: float = Field(
|
|
default=0.5, ge=0, le=1, description="Ratio of tokens to keep"
|
|
)
|
|
hidden_ratio: float = Field(
|
|
default=2.0, gt=0, description="MLP hidden dim as multiple of compression_dim"
|
|
)
|
|
dropout_rate: float = Field(
|
|
default=0.1, ge=0, le=1, description="Dropout probability"
|
|
)
|
|
use_residual: bool = True
|
|
device: str = "auto"
|
|
|
|
|
|
class VisualizationConfig(BaseModel):
|
|
"""Configuration for visualization settings."""
|
|
|
|
model_config = ConfigDict(extra="ignore")
|
|
|
|
plot_theme: str = "plotly_white"
|
|
color_scale: str = "viridis"
|
|
point_size: int = Field(default=8, gt=0)
|
|
fig_width: int = Field(default=900, gt=0)
|
|
fig_height: int = Field(default=600, gt=0)
|
|
|
|
|
|
class OutputConfig(BaseModel):
|
|
"""Configuration for output settings."""
|
|
|
|
model_config = ConfigDict(extra="ignore")
|
|
|
|
directory: Path = Path(__file__).parent.parent.parent / "outputs"
|
|
html_self_contained: bool = True
|
|
png_scale: int = Field(default=2, gt=0)
|
|
|
|
@field_validator("directory", mode="after")
|
|
def convert_to_absolute(cls, v: Path) -> Path:
|
|
"""
|
|
Converts the path to an absolute path relative to the current working directory.
|
|
This works even if the path doesn't exist on disk.
|
|
"""
|
|
if v.is_absolute():
|
|
return v
|
|
return Path(__file__).parent.parent.parent / v
|
|
|
|
|
|
class FeatureCompressorConfig(BaseModel):
|
|
"""Root configuration for the feature compressor."""
|
|
|
|
model_config = ConfigDict(extra="ignore")
|
|
|
|
model: ModelConfig
|
|
visualization: VisualizationConfig
|
|
output: OutputConfig
|