type hints and favicons

Author: FBruzzesi

pyproject.toml

@@ -5,34 +5,55 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "fastapi[standard]>=0.115.5",
-    "geopy>=2.4.1",
-    "httpx>=0.27.2",
-    "numpy<2",
+    "numpy==1.26.4",
     "pillow==10.4.0",
-    "sentence-transformers>=3.3.1",
-    "sqlmodel>=0.0.22",
-    "streamlit>=1.40.1",
+]
+
+[project.optional-dependencies]
+frontend = [
+    "geopy==2.4.1",
+    "httpx==0.27.2",
+    "streamlit==1.40.1",
+]
+
+backend = [
+    "fastapi[standard]==0.115.5",
+    "sqlmodel==0.0.22",
+    "sentence-transformers==3.3.1",
+    "torch==2.5.1+cpu",
 ]
 
 [[tool.uv.index]]
 name = "pytorch-cpu"
 url = "https://download.pytorch.org/whl/cpu"
 explicit = true
 
+[tool.uv.sources]
+torch = [{ index = "pytorch-cpu"}]
+
 [project.scripts]
 moin-moin = "moin_moin:main"
 
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
-[dependency-groups]
-dev = [
-    "fastbook>=0.0.29",
-    "ipdb>=0.13.13",
-    "ipyplot>=1.1.2",
-    "ipython>=8.29.0",
-    "jupyterlab>=4.3.1",
-    "jupyterlab-vim>=4.1.4",
+
+[tool.ruff]
+line-length = 120
+target-version = "py312"
+fix = true
+
+lint.select=["ALL"]
+
+lint.ignore = [
+  "COM812",
+  "ISC001",
+  "D203",
+  "D211",
+  "D213",
 ]
+
+[tool.ruff.lint.isort]
+force-single-line = true
+required-imports = ["from __future__ import annotations"]

src/moin_moin/__init__.py

@@ -1,2 +1,5 @@
-def main() -> None:
-    print("Hello from moin-moin!")
+"""moin-moin app.
+
+backend: contains fastapi endpoint as well as machine learning and database components.
+frontend: streamlit app to render webapp.
+"""

src/moin_moin/backend/__init__.py

@@ -0,0 +1 @@
+"""Backend: API and DB for app."""

src/moin_moin/backend/_db.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from sqlmodel import Field
+from sqlmodel import SQLModel
+
+
+class UserUploadData(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    image: bytes
+    latitude: float
+    longitude: float
+    notes: str
+    tags: str
+
+
+class Prediction(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    record_id: int | None = Field(default=None, foreign_key="useruploaddata.id")
+    prediction: str
+
+
+class PublicRecord(SQLModel):
+    latitude: float
+    longitude: float
+    notes: str
+    tags: str
+    prediction: str

src/moin_moin/backend/_ml.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Self
+
+from sentence_transformers import SentenceTransformer
+
+if TYPE_CHECKING:
+    from PIL import Image
+
+
+class ClipModel:
+    """Wrapper of SentenceTransformer, clip-ViT-B-32 model.
+
+    Process is to store embeddings of categories we can match images with, and then at
+    prediction time to return the category closest to the given image.
+    """
+
+    def __init__(self: Self, text_options: dict[str, str]) -> None:
+        self.model = SentenceTransformer("clip-ViT-B-32")
+        self.model.eval()
+        self.labels = list(text_options.keys())
+        self.text_embedding = self.model.encode(list(text_options.values()))
+
+    def predict(self: Self, image: Image) -> str:
+        """Embed the image and return the closest text embedding."""
+        img_emb = self.model.encode(image)
+        similarity_array = self.model.similarity(img_emb, self.text_embedding)
+        return self.labels[similarity_array.argmax().item()]

src/moin_moin/backend/api.py

@@ -1,25 +1,42 @@
+"""Module that builds FastAPI application with endpoints to save, predict and gather data."""
+
+from __future__ import annotations
+
 from contextlib import asynccontextmanager
+from io import BytesIO
+from typing import TYPE_CHECKING
 from typing import Annotated
+from typing import Final
 
-from fastapi import FastAPI, File, UploadFile, Form, Depends
+from fastapi import Depends
+from fastapi import FastAPI
+from fastapi import File
+from fastapi import Form
+from fastapi import UploadFile
 from PIL import Image
-from io import BytesIO
-from sqlmodel import Session, select
-from moin_moin.backend.ml import ClipModel
-from moin_moin.backend.db import Record, Prediction, engine, PublicRecord
+from sqlmodel import Session
+from sqlmodel import SQLModel
+from sqlmodel import create_engine
+from sqlmodel import select
 
+from moin_moin.backend._db import Prediction
+from moin_moin.backend._db import PublicRecord
+from moin_moin.backend._db import UserUploadData
+from moin_moin.backend._ml import ClipModel
 
-ML_MODEL = {}
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+    from collections.abc import Generator
 
+    from sqlmodel import Engine
 
-def get_session():
-    with Session(engine) as session:
-        yield session
 
+ML_MODEL: dict[str, ClipModel] = {}
 
-SessionDep = Annotated[Session, Depends(get_session)]
+DB_NAME: Final[str] = "sqlite:///moin-moin.db"
+ENGINE: Final[Engine] = create_engine(DB_NAME)
 
-institutions = {
+INSTITUTIONS = {
     "Police Department": "The police deals with crime and violence related topics.",
     "Fire Department": "The fire department deals with fire and other emergency situations.",
     "Hospital": "The hospital deals with health and medical related topics.",
@@ -28,42 +45,48 @@ def get_session():
 }
 
 
+def get_session() -> Generator[Session, None, None]:
+    """Yield SQLModel session."""
+    with Session(ENGINE) as session:
+        yield session
+
+
+SessionDep = Annotated[Session, Depends(get_session)]
+
 
 @asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Load the ML model
-    ML_MODEL["similarity_model"] = ClipModel(text_options=institutions)
+async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:  # noqa: ARG001
+    """FastAPI Lifespan: loads and instantiate the model at startup, delete at shutdown."""
+    SQLModel.metadata.create_all(ENGINE)
+    ML_MODEL["similarity_model"] = ClipModel(text_options=INSTITUTIONS)
+
     yield
-    # Clean up the ML models and release the resources
+
     ML_MODEL.clear()
 
 
 app = FastAPI(lifespan=lifespan)
 
 
-def _predict(img_array: Image):
-    """This function decoupled from the API exists for debugging purposes."""
-    return ML_MODEL["similarity_model"].predict(img_array)
-
-
 @app.get("/")
-def root():
+def root() -> dict[str, str]:
     """Root endpoint."""
     return {"message": "Health"}
 
 
 @app.post("/save")
-async def save(
+async def save(  # noqa: PLR0913
     session: SessionDep,
-    latitude: float = Form(...),
-    longitude: float = Form(...),
-    notes: str = Form(...),
-    tags: str = Form(...),
-    image_bytes: UploadFile = File(...),
-):
+    latitude: Annotated[float, Form()],
+    longitude: Annotated[float, Form()],
+    notes: Annotated[str, Form()],
+    tags: Annotated[str, Form()],
+    image_bytes: Annotated[UploadFile, File()],
+) -> dict[str, int | None]:
+    """Save the input data into database records into user upload data table."""
     image_bytes = await image_bytes.read()
-    record = Record(
-        image=image_bytes,
+    record = UserUploadData(
+        image=image_bytes,  # type: ignore[assignment]
         latitude=latitude,
         longitude=longitude,
         notes=notes,
@@ -79,12 +102,15 @@ async def save(
 
 @app.post("/predict")
 async def predict(
-    session: SessionDep, record_id: int = Form(), file: UploadFile = File(...)
-):
+    session: SessionDep,
+    record_id: Annotated[int, Form()],
+    file: Annotated[UploadFile, File(...)],
+) -> dict[str, str]:
+    """Predict image category via ML model."""
     file_bytes = await file.read()
     buffer = BytesIO(file_bytes)
     image = Image.open(buffer)
-    prediction = _predict(image)
+    prediction = ML_MODEL["similarity_model"].predict(image)
 
     pred_record = Prediction(record_id=record_id, prediction=prediction)
 
@@ -95,13 +121,14 @@ async def predict(
 
 
 @app.get("/load-records", response_model=list[PublicRecord])
-async def load_records(session: SessionDep):
+async def load_records(session: SessionDep) -> list[PublicRecord]:
+    """Load all the records with their predictions, ignores associated image."""
     statement = select(
-        Record.latitude,
-        Record.longitude,
-        Record.notes,
-        Record.tags,
+        UserUploadData.latitude,
+        UserUploadData.longitude,
+        UserUploadData.notes,
+        UserUploadData.tags,
         Prediction.prediction,
     ).join(Prediction)
     records = session.exec(statement).all()
-    return [PublicRecord(**row._mapping) for row in records]
+    return [PublicRecord(**row._mapping) for row in records]  # noqa: SLF001

src/moin_moin/backend/db.py

@@ -0,0 +1 @@
+"""Frontend: Streamlit app."""

src/moin_moin/backend/ml.py

@@ -1,7 +1,10 @@
-HOST = "http://127.0.0.1"
-PORT = 8000
+from __future__ import annotations
 
-INSTITUTION_MAPPING = {
+from typing import Final
+
+HOST: Final[str] = "http://127.0.0.1"
+PORT: Final[int] = 8081
+INSTITUTION_MAPPING: Final[dict[str, str]] = {
     "Police Department": "#48b5a5",
     "Fire Department": "#7b64ab",
     "Hospital": "#20be64",