Initial commit
This commit is contained in:
tanexc
@@ -0,0 +1,119 @@
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from onnxruntime.quantization import (
|
||||
CalibrationDataReader,
|
||||
QuantFormat,
|
||||
QuantType,
|
||||
quantize_static,
|
||||
)
|
||||
from onnxruntime.tools.convert_onnx_models_to_ort import (
|
||||
OptimizationStyle,
|
||||
convert_onnx_models_to_ort,
|
||||
)
|
||||
from PIL import Image
|
||||
from ultralytics import YOLO
|
||||
|
||||
from app.domain.ModelConverter import ModelConverter
|
||||
from app.infrastructure.constants import BASE_DIR
|
||||
|
||||
|
||||
IMAGES_DIR = BASE_DIR / "storage" / "images"
|
||||
|
||||
|
||||
class ModelConverterImpl(ModelConverter):
|
||||
def convert(self, model_pt_path: Path, results_dir: Path) -> Path:
|
||||
model_pt_path = Path(model_pt_path)
|
||||
results_dir = Path(results_dir)
|
||||
results_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
result_path = results_dir / f"{model_pt_path.stem}.ort"
|
||||
if result_path.exists():
|
||||
result_path.unlink()
|
||||
|
||||
with TemporaryDirectory(dir=results_dir) as temp_dir_name:
|
||||
temp_dir = Path(temp_dir_name)
|
||||
fp32_onnx_path = temp_dir / f"{model_pt_path.stem}_fp32.onnx"
|
||||
int8_onnx_path = temp_dir / f"{model_pt_path.stem}_int8_static.onnx"
|
||||
|
||||
model = YOLO(str(model_pt_path), task="detect")
|
||||
exported_path = Path(
|
||||
model.export(
|
||||
format="onnx",
|
||||
imgsz=640,
|
||||
opset=20,
|
||||
dynamic=False,
|
||||
simplify=False,
|
||||
half=False,
|
||||
)
|
||||
).resolve()
|
||||
shutil.move(str(exported_path), str(fp32_onnx_path))
|
||||
|
||||
calibration_reader = ImageCalibrationDataReader(fp32_onnx_path, IMAGES_DIR)
|
||||
quantize_static(
|
||||
model_input=str(fp32_onnx_path),
|
||||
model_output=str(int8_onnx_path),
|
||||
calibration_data_reader=calibration_reader,
|
||||
quant_format=QuantFormat.QDQ,
|
||||
activation_type=QuantType.QUInt8,
|
||||
weight_type=QuantType.QInt8,
|
||||
op_types_to_quantize=["Conv"],
|
||||
)
|
||||
|
||||
convert_onnx_models_to_ort(
|
||||
int8_onnx_path,
|
||||
output_dir=temp_dir,
|
||||
optimization_styles=[OptimizationStyle.Fixed],
|
||||
)
|
||||
generated_ort_path = temp_dir / f"{int8_onnx_path.stem}.ort"
|
||||
generated_ort_path.replace(result_path)
|
||||
|
||||
return result_path
|
||||
|
||||
|
||||
class ImageCalibrationDataReader(CalibrationDataReader):
|
||||
def __init__(self, model_path: Path, images_dir: Path, limit: int = 32):
|
||||
self.session = ort.InferenceSession(
|
||||
str(model_path), providers=["CPUExecutionProvider"]
|
||||
)
|
||||
self.input_name = self.session.get_inputs()[0].name
|
||||
self.input_shape = self.session.get_inputs()[0].shape
|
||||
self.image_paths = [
|
||||
path
|
||||
for path in sorted(images_dir.glob("*"))
|
||||
if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp"}
|
||||
][:limit]
|
||||
self.index = 0
|
||||
|
||||
if not self.image_paths:
|
||||
raise FileNotFoundError(f"calibration images not found in {images_dir}")
|
||||
|
||||
def get_next(self) -> dict[str, np.ndarray] | None:
|
||||
if self.index >= len(self.image_paths):
|
||||
return None
|
||||
|
||||
image_path = self.image_paths[self.index]
|
||||
self.index += 1
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
|
||||
height = int(self.input_shape[2]) if isinstance(self.input_shape[2], int) else 640
|
||||
width = int(self.input_shape[3]) if isinstance(self.input_shape[3], int) else 640
|
||||
scale = min(width / image.width, height / image.height)
|
||||
resized_width = int(round(image.width * scale))
|
||||
resized_height = int(round(image.height * scale))
|
||||
resized_image = image.resize((resized_width, resized_height), Image.BILINEAR)
|
||||
|
||||
padded_image = Image.new("RGB", (width, height), (114, 114, 114))
|
||||
left = (width - resized_width) // 2
|
||||
top = (height - resized_height) // 2
|
||||
padded_image.paste(resized_image, (left, top))
|
||||
|
||||
image_array = np.asarray(padded_image, dtype=np.float32) / 255.0
|
||||
image_array = np.transpose(image_array, (2, 0, 1))
|
||||
return {self.input_name: np.expand_dims(image_array, axis=0)}
|
||||
|
||||
def rewind(self) -> None:
|
||||
self.index = 0
|
||||
@@ -0,0 +1,47 @@
|
||||
from sqlalchemy import select, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import Session
|
||||
from app.domain.Exception import DuplicateModelVersionError
|
||||
from app.domain.Model import Model
|
||||
from app.domain.ModelRepository import ModelRepository
|
||||
from app.infrastructure.database import ModelRecord
|
||||
|
||||
|
||||
class ModelRepositoryImpl(ModelRepository):
|
||||
def __init__(self, session: Session) -> None:
|
||||
self._session = session
|
||||
|
||||
def get_latest(self) -> Model | None:
|
||||
stmt = (
|
||||
select(ModelRecord)
|
||||
.order_by(ModelRecord.is_latest.desc(), ModelRecord.created_at.desc())
|
||||
.limit(1)
|
||||
)
|
||||
record = self._session.scalar(stmt)
|
||||
return record.to_domain() if record else None
|
||||
|
||||
def get_by_version(self, version: str) -> Model | None:
|
||||
stmt = select(ModelRecord).where(ModelRecord.version == version)
|
||||
record = self._session.scalar(stmt)
|
||||
return record.to_domain() if record else None
|
||||
|
||||
def add(self, model: Model) -> Model:
|
||||
record = ModelRecord(
|
||||
version=model.version,
|
||||
filename=model.filename,
|
||||
file_path=model.file_path,
|
||||
created_at=model.created_at,
|
||||
is_latest=model.is_latest,
|
||||
)
|
||||
try:
|
||||
self._session.execute(update(ModelRecord).values(is_latest=False))
|
||||
self._session.add(record)
|
||||
self._session.commit()
|
||||
self._session.refresh(record)
|
||||
return record.to_domain()
|
||||
except IntegrityError as exc:
|
||||
self._session.rollback()
|
||||
raise DuplicateModelVersionError(
|
||||
f"Model version '{model.version}' already exists."
|
||||
) from exc
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
from app.domain.Exception import InvalidModelFileError, ModelFileMissingError
|
||||
from app.domain.ModelStorage import ModelStorage
|
||||
from app.infrastructure.constants import TMP_DIR, MODEL_DIR
|
||||
|
||||
|
||||
class ModelStorageImpl(ModelStorage):
|
||||
@property
|
||||
def models_dir(self) -> Path:
|
||||
return MODEL_DIR
|
||||
|
||||
def ensure_directories(self) -> None:
|
||||
MODEL_DIR.mkdir(parents=True, exist_ok=True)
|
||||
TMP_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def validate_upload_filename(self, filename: str | None) -> str:
|
||||
if not filename:
|
||||
raise InvalidModelFileError("Uploaded file must have a filename.")
|
||||
|
||||
safe_name = Path(filename).name
|
||||
if Path(safe_name).suffix.lower() != ".pt":
|
||||
raise InvalidModelFileError("Only .pt model files are supported.")
|
||||
return safe_name
|
||||
|
||||
def save_upload_to_tmp(self, source_file, filename: str) -> Path:
|
||||
tmp_path = TMP_DIR / f"{uuid4().hex}_{filename}"
|
||||
with tmp_path.open("wb") as buffer:
|
||||
shutil.copyfileobj(source_file, buffer)
|
||||
return tmp_path
|
||||
|
||||
def build_model_filename(self, version: str) -> str:
|
||||
return f"model_v{version}.ort"
|
||||
|
||||
def prepare_conversion_input(self, tmp_path: Path, model_filename: str) -> Path:
|
||||
conversion_input_path = TMP_DIR / f"{uuid4().hex}_{Path(model_filename).stem}.pt"
|
||||
tmp_path.rename(conversion_input_path)
|
||||
return conversion_input_path
|
||||
|
||||
def finalize_converted_model(self, converted_path: Path, model_filename: str) -> Path:
|
||||
target_path = MODEL_DIR / model_filename
|
||||
if converted_path.resolve() != target_path.resolve():
|
||||
converted_path.replace(target_path)
|
||||
return target_path
|
||||
|
||||
def ensure_model_file_exists(self, file_path: str) -> Path:
|
||||
resolved_path = Path(file_path)
|
||||
if not resolved_path.exists():
|
||||
raise ModelFileMissingError("Latest model file is missing on disk.")
|
||||
return resolved_path
|
||||
|
||||
def delete_file_if_exists(self, file_path: Path | None) -> None:
|
||||
if file_path and file_path.exists():
|
||||
file_path.unlink()
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from app.infrastructure.constants import DATABASE_URL
|
||||
from app.infrastructure.ModelConverterImpl import ModelConverterImpl
|
||||
from app.infrastructure.database import Base, create_engine_and_session
|
||||
from app.infrastructure.ModelStorageImpl import ModelStorageImpl
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Container:
|
||||
storage: ModelStorageImpl
|
||||
converter: ModelConverterImpl
|
||||
engine: object
|
||||
session_factory: object
|
||||
|
||||
|
||||
_container: Container | None = None
|
||||
|
||||
|
||||
def get_container() -> Container:
|
||||
global _container
|
||||
if _container is None:
|
||||
storage = ModelStorageImpl()
|
||||
engine, session_factory = create_engine_and_session(DATABASE_URL)
|
||||
_container = Container(
|
||||
storage=storage,
|
||||
converter=ModelConverterImpl(),
|
||||
engine=engine,
|
||||
session_factory=session_factory,
|
||||
)
|
||||
return _container
|
||||
|
||||
|
||||
def initialize_application() -> None:
|
||||
container = get_container()
|
||||
container.storage.ensure_directories()
|
||||
engine_path = container.engine.url.database
|
||||
if engine_path:
|
||||
Path(engine_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
Base.metadata.create_all(bind=container.engine)
|
||||
@@ -0,0 +1,13 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
BASE_DIR = Path(__file__).resolve().parent.parent.parent
|
||||
MODEL_DIR = Path(os.getenv("MODELS_DIR", str(BASE_DIR / "storage" / "models")))
|
||||
TMP_DIR = Path(os.getenv("TMP_DIR", str(BASE_DIR / "tmp")))
|
||||
DATABASE_PATH = BASE_DIR / "db" / "app.db"
|
||||
DATABASE_URL = os.getenv("DATABASE_URL", f"sqlite:///{DATABASE_PATH}")
|
||||
API_KEY_HEADER_NAME = "X-API-Key"
|
||||
API_KEYS = tuple(
|
||||
key.strip() for key in os.getenv("API_KEYS", "").split(",") if key.strip()
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
from datetime import datetime, timezone
|
||||
from sqlalchemy import Boolean, DateTime, Integer, String, create_engine
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
|
||||
from app.domain.Model import Model
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
class ModelRecord(Base):
|
||||
__tablename__ = "models"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
|
||||
version: Mapped[str] = mapped_column(String(64), unique=True, nullable=False, index=True)
|
||||
filename: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
|
||||
file_path: Mapped[str] = mapped_column(String(1024), nullable=False)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
)
|
||||
is_latest: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, index=True)
|
||||
|
||||
def to_domain(self) -> Model:
|
||||
return Model(
|
||||
id=self.id,
|
||||
version=self.version,
|
||||
filename=self.filename,
|
||||
file_path=self.file_path,
|
||||
created_at=self.created_at,
|
||||
is_latest=self.is_latest,
|
||||
)
|
||||
|
||||
|
||||
def create_engine_and_session(database_url: str):
|
||||
engine = create_engine(
|
||||
database_url,
|
||||
connect_args={"check_same_thread": False} if database_url.startswith("sqlite") else {},
|
||||
)
|
||||
session_factory = sessionmaker(bind=engine, autocommit=False, autoflush=False)
|
||||
return engine, session_factory
|
||||
Reference in New Issue
Block a user