mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-03-12 12:25:32 +08:00
52 lines
1.3 KiB
Python
52 lines
1.3 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=512, gt=0, description="Output feature dimension"
|
|
)
|
|
device: str = "auto"
|
|
|
|
|
|
class OutputConfig(BaseModel):
|
|
"""Configuration for output settings."""
|
|
|
|
model_config = ConfigDict(extra="ignore")
|
|
|
|
directory: Path = Path(__file__).parent.parent.parent / "outputs"
|
|
|
|
@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 Config(BaseModel):
|
|
"""Root configuration for the feature compressor."""
|
|
|
|
model_config = ConfigDict(extra="ignore")
|
|
|
|
model: ModelConfig
|
|
output: OutputConfig
|