90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
import json
|
|
from abc import abstractmethod
|
|
from functools import cached_property
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from numpy.typing import NDArray
|
|
from PIL import Image
|
|
|
|
from immich_ml.config import log
|
|
from immich_ml.models.base import InferenceModel
|
|
from immich_ml.models.transforms import (
|
|
crop_pil,
|
|
decode_pil,
|
|
get_pil_resampling,
|
|
normalize,
|
|
resize_pil,
|
|
serialize_np_array,
|
|
to_numpy,
|
|
)
|
|
from immich_ml.schemas import ModelGraph, ModelInput, ModelSession, ModelSource, ModelTask, ModelType, VisualOptions
|
|
|
|
|
|
class BaseCLIPVisualEncoder(InferenceModel[VisualOptions]):
|
|
depends = []
|
|
identity = (ModelType.VISUAL, ModelTask.SEARCH)
|
|
|
|
def _predict(self, inputs: Image.Image | bytes, options: VisualOptions) -> str:
|
|
image = decode_pil(inputs)
|
|
session = self.session.for_shape(self.shape_policy.dims[0])
|
|
res: NDArray[np.float32] = session.run(None, self.transform(session, image))[0][0]
|
|
return serialize_np_array(res)
|
|
|
|
@abstractmethod
|
|
def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput:
|
|
pass
|
|
|
|
@property
|
|
def model_cfg_path(self) -> Path:
|
|
return self.cache_dir / "config.json"
|
|
|
|
@property
|
|
def preprocess_cfg_path(self) -> Path:
|
|
return self.model_dir / "preprocess_cfg.json"
|
|
|
|
@cached_property
|
|
def model_cfg(self) -> dict[str, Any]:
|
|
log.debug(f"Loading model config for CLIP model '{self.model_name}'")
|
|
model_cfg: dict[str, Any] = json.load(self.model_cfg_path.open(encoding="utf-8"))
|
|
log.debug(f"Loaded model config for CLIP model '{self.model_name}'")
|
|
return model_cfg
|
|
|
|
@cached_property
|
|
def preprocess_cfg(self) -> dict[str, Any]:
|
|
log.debug(f"Loading visual preprocessing config for CLIP model '{self.model_name}'")
|
|
preprocess_cfg: dict[str, Any] = json.load(self.preprocess_cfg_path.open(encoding="utf-8"))
|
|
log.debug(f"Loaded visual preprocessing config for CLIP model '{self.model_name}'")
|
|
return preprocess_cfg
|
|
|
|
|
|
class OpenClipVisualEncoder(BaseCLIPVisualEncoder):
|
|
sources = (ModelSource.OPENCLIP, ModelSource.MCLIP)
|
|
|
|
def _load(self) -> ModelSession:
|
|
size: list[int] | int = self.preprocess_cfg["size"]
|
|
self.size = size[0] if isinstance(size, list) else size
|
|
|
|
self.resampling = get_pil_resampling(self.preprocess_cfg["interpolation"])
|
|
self.mean = np.array(self.preprocess_cfg["mean"], dtype=np.float32)
|
|
self.std = np.array(self.preprocess_cfg["std"], dtype=np.float32)
|
|
|
|
return super()._load()
|
|
|
|
def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput:
|
|
image = self._resize(image)
|
|
if session.normalizes_input:
|
|
rgb: NDArray[np.uint8] = np.asarray(image if image.mode == "RGB" else image.convert("RGB"))
|
|
return {"image": rgb[None]}
|
|
return {"image": normalize(to_numpy(image), self.mean, self.std).transpose(2, 0, 1)[None]}
|
|
|
|
def _resize(self, image: Image.Image) -> Image.Image:
|
|
match self.preprocess_cfg.get("resize_mode", "shortest"): # open_clip's default for older configs
|
|
case "squash":
|
|
return image.resize((self.size, self.size), resample=self.resampling)
|
|
case "shortest":
|
|
return crop_pil(resize_pil(image, self.size, self.resampling), self.size)
|
|
case mode:
|
|
raise ValueError(f"Unsupported resize_mode {mode!r} in {self.preprocess_cfg_path}")
|