feat(routes): reorganize routes and add embedding functionality for t…

…ext and image processing

.gitignore

@@ -1,8 +1,6 @@
-# FastAPI dev/build outputs
+# Python dependencies
 __pycache__
 .ruff_cache
-
-# Python dependencies
 .venv
 
 # Logs

Dockerfile

@@ -1,21 +1,20 @@
 FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim AS builder
 
-# Enable bytecode compilation
 ENV UV_COMPILE_BYTECODE=1
-# Copy from the cache instead of linking since it's a mounted volume
 ENV UV_LINK_MODE=copy
 
+WORKDIR /app
+
 RUN --mount=type=cache,target=/root/.cache/uv \
   --mount=type=bind,source=uv.lock,target=uv.lock \
   --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
   uv sync --frozen --no-install-project --no-dev
 
-ADD . /app
+COPY .python-version pyproject.toml uv.lock ./
 
-WORKDIR /app
+RUN --mount=type=cache,target=/root/.cache/uv uv sync --frozen --no-dev
 
-RUN --mount=type=cache,target=/root/.cache/uv \
-  uv sync --frozen --no-dev
+COPY . .
 
 FROM python:3.12-slim AS runner
 
@@ -25,12 +24,10 @@ ARG BUILD_TIME
 ENV PYTHON_ENV=production
 ENV PATH="/app/.venv/bin:$PATH"
 
-COPY --from=builder /app/.venv /app/.venv
-COPY --from=builder /app/assets /app/assets
-COPY --from=builder /app/server /app/server
-
 WORKDIR /app
 
+COPY --from=builder /app/.venv /app/assets /app/server ./
+
 EXPOSE 8000
 
 CMD ["fastapi", "run", "server/main.py", "--host", "0.0.0.0", "--port", "8000"]

pyproject.toml

@@ -14,13 +14,16 @@ dependencies = [
     "pydantic>=2.10.4",
     "scipy>=1.14.1",
     "nanoid>=2.0.0",
+    "transformers[torch]>=4.47.1",
+    "timm>=1.0.12",
+    "einops>=0.8.0",
 ]
 
 [dependency-groups]
 dev = [
     "changelog-gen>=0.13.5",
     "pytest>=8.3.4",
-    "ruff>=0.8.3",
+    "ruff>=0.8.4",
     "taskipy>=1.14.1",
 ]
 

server/main.py

@@ -1,11 +1,13 @@
 from fastapi import FastAPI
 
 # from server.dependencies import get_query_token
-from server.routes import health, detect, omr, search
 
-app = FastAPI()  # dependencies=[Depends(get_query_token)])
+from server.routes import health, image, text
+
+# dependencies=[Depends(get_query_token)])
+app = FastAPI()
 
 app.include_router(health.router)
-app.include_router(detect.router)
-app.include_router(search.router)
-app.include_router(omr.router)
+app.include_router(image.detect.router)
+app.include_router(image.embedding.router)
+app.include_router(text.embedding.router)

server/routes/image/__init__.py

@@ -0,0 +1,5 @@
+from . import detect
+from . import embedding
+
+
+__all__ = [detect, embedding]

server/routes/detect.py → server/routes/image/detect.py

@@ -1,10 +1,12 @@
 import json
 from typing import List, Tuple
-from fastapi import APIRouter, HTTPException
+from fastapi import HTTPException
+from pydantic import BaseModel
 import httpx
 from nanoid import generate
 import numpy as np
-from pydantic import BaseModel
+
+from .router import router
 
 from server.core.config import config
 from server.utils.base64ToArray import base64_to_array
@@ -83,31 +85,24 @@ def predict(image: np.ndarray):
 
     try:
         detections = httpx.post(
-            f"{config.tensorflow_api_url}/v1/models/detector:predict",
+            f"{config.tensorflow_api_url}/v1/models/custom-detector:predict",
             data=data,
         )
 
         return postprocess(detections.json()["predictions"][0], dim)
     except Exception as error:
-        print("Failed request Tensorflow Serving /detector:predict", error)
+        print("Failed request Tensorflow Serving /custom-detector:predict", error)
         raise HTTPException(
             status_code=500,
-            detail="Failed request Tensorflow Serving /detector:predict",
+            detail="Failed request Tensorflow Serving /custom-detector:predict",
         )
 
 
-router = APIRouter(
-    prefix="/detect",
-    tags=["detect"],
-    responses={404: {"description": "Not found"}},
-)
-
-
 class RequestBody(BaseModel):
     image: str
 
 
-@router.post("/")
+@router.post("/detect")
 async def detect(request: RequestBody):
     try:
         id = generate()

server/routes/search.py → server/routes/image/embedding.py

@@ -2,12 +2,14 @@
 import json
 import math
 from typing import List, Tuple
-from fastapi import APIRouter, HTTPException
+from fastapi import HTTPException
 from pydantic import BaseModel
 import httpx
 import numpy as np
 from PIL import Image
 
+from .router import router
+
 from server.core.config import config
 from server.utils.convertBoxFormat import convert_box_format
 from server.utils.imageLoad import image_load
@@ -38,18 +40,18 @@ async def predict(image: np.ndarray, boxes: list):
 
     try:
         detections = httpx.post(
-            f"{config.tensorflow_api_url}/v1/models/feature_extractor:predict",
+            f"{config.tensorflow_api_url}/v1/models/custom-embedder:predict",
             data=data,
         )
 
         embeddings = postprocess(detections.json()["predictions"])
 
         return embeddings
     except Exception as error:
-        print("Failed request Tensorflow Serving /feature_extractor:predict", error)
+        print("Failed request Tensorflow Serving /custom-embedder:predict", error)
         raise HTTPException(
             status_code=500,
-            detail="Failed request Tensorflow Serving /feature_extractor:predict",
+            detail="Failed request Tensorflow Serving /custom-embedder:predict",
         )
 
 
@@ -83,34 +85,26 @@ def process_box(box):
     return np.array(cropped_images)
 
 
-router = APIRouter(
-    prefix="/search",
-    tags=["search"],
-    responses={404: {"description": "Not found"}},
-)
-
-
 class Object(BaseModel):
     box: List[float]
     confidence: float
-    category: str
 
 
 class RequestBody(BaseModel):
     id: str
     objects: List[Object]
 
 
-@router.post("/")
-async def search(request: RequestBody):
+@router.post("/embedding")
+async def get_embedding(request: RequestBody):
     try:
         image = image_load(request.id)
 
-        searches = await predict(
+        embeddings = await predict(
             image, list(map(lambda object: object.box, request.objects))
         )
 
-        return {"id": request.id, "searches": searches}
+        return {"id": request.id, "embeddings": embeddings}
     except HTTPException as error:
         print("API search POST", error)
         raise error

server/routes/omr.py → server/routes/image/omr.py

@@ -1,7 +1,9 @@
-from fastapi import APIRouter, HTTPException
+from fastapi import HTTPException
 from nanoid import generate
 from pydantic import BaseModel
 
+from .router import router
+
 from server.utils.base64ToArray import base64_to_array
 from server.utils.imageSave import image_save
 from server.utils.omrAlignCrop import align_crop
@@ -11,18 +13,12 @@
 from server.utils.omrExtractData import extract_data
 from server.utils.omrHighlights import get_highlights
 
-router = APIRouter(
-    prefix="/omr",
-    tags=["omr"],
-    responses={404: {"description": "Not found"}},
-)
-
 
 class RequestBody(BaseModel):
     image: str
 
 
-@router.post("/")
+@router.post("/omr")
 async def omr(request: RequestBody):
     try:
         id = generate()

server/routes/image/router.py

@@ -0,0 +1,8 @@
+from fastapi import APIRouter
+
+
+router = APIRouter(
+    prefix="/image",
+    tags=["image"],
+    responses={404: {"description": "Not found"}},
+)

server/routes/text/__init__.py

@@ -0,0 +1,3 @@
+from . import embedding
+
+__all__ = [embedding]

server/routes/text/embedding.py

@@ -0,0 +1,33 @@
+from typing import List
+from fastapi import HTTPException
+from pydantic import BaseModel
+from transformers import AutoModel
+
+from .router import router
+
+model = AutoModel.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True)
+
+
+class RequestBody(BaseModel):
+    id: str
+    queries: List[str]
+
+
+@router.post("/embedding")
+async def get_embedding(request: RequestBody):
+    try:
+        truncate_dim = 512
+
+        embeddings = model.encode_text(request.queries, truncate_dim=truncate_dim)
+
+        return {
+            "id": request.id,
+            "embeddings": [[float(x) for x in embedding] for embedding in embeddings],
+        }
+        return
+    except HTTPException as error:
+        print("API similarity POST", error)
+        raise error
+    except Exception as error:
+        print("API similarity POST", error)
+        raise HTTPException(status_code=500, detail="Some Unknown Error Found")

server/routes/text/router.py

@@ -0,0 +1,8 @@
+from fastapi import APIRouter
+
+
+router = APIRouter(
+    prefix="/text",
+    tags=["text"],
+    responses={404: {"description": "Not found"}},
+)