feat: add save tags endpoint for image URLs

Author: herzo175

tagger/api/schema/tags.py

@@ -25,3 +25,9 @@ class Tags(BaseModel):
 
 class TagsResponse(BaseModel):
     tags: List[Tags]
+
+
+class SaveTagsRequest(BaseModel):
+    category: str
+    image: Image
+    tags: List[Tags]

tagger/api/v1/tags.py

@@ -1,11 +1,39 @@
 from fastapi import APIRouter
 
-from tagger.api.schema.tags import TagsRequest, TagsResponse
-from tagger.core.tags import generate_tags
+from tagger.api.schema.tags import SaveTagsRequest, TagsRequest, TagsResponse
+from tagger.config.models import VISION_EMBEDDING_MODEL
+from tagger.core.tags import (
+    download_image_url,
+    generate_tags,
+    save_tag_embedding,
+    resize_image,
+)
 
 router = APIRouter(prefix="/tags")
 
 
 @router.post("/", response_model=TagsResponse)
 async def create_tags(tag: TagsRequest):
     return generate_tags(tag)
+
+
+@router.post("/save", response_model=TagsResponse)
+async def save_tags(tag: SaveTagsRequest):
+    """
+    Save generated tags for an image to the database.
+    """
+
+    # Download image from url
+    base64_image = resize_image(download_image_url(tag.image.url))
+
+    # Generate image embedding
+    image_embedding_value = VISION_EMBEDDING_MODEL.image_embedding([base64_image])[0]
+
+    # Save image embedding + tags to database
+    save_tag_embedding(
+        category=tag.category,
+        image_url=tag.image.url,
+        image_embeddings=image_embedding_value,
+        coordinates=tag.image.coordinates,
+        tags=tag.tags,
+    )

tagger/core/tags.py

@@ -13,7 +13,7 @@
 import requests
 import boto3
 
-from tagger.api.schema.tags import Tags, TagsRequest, TagsResponse
+from tagger.api.schema.tags import Coordinates, Tags, TagsRequest, TagsResponse
 from tagger.config.models import JSON_OUTPUT_MODEL, VISION_EMBEDDING_MODEL, VISION_MODEL
 from tagger.config.db import TAGGING_DB_ENGINE
 from tagger.config.storage import S3_CLIENT
@@ -218,6 +218,26 @@ def get_similar_images(
         ]
 
 
+def save_tag_embedding(
+    category: str,
+    image_url: str,
+    image_embeddings: List[float],
+    coordinates: Coordinates,
+    tags: List[Tags],
+):
+    with Session(TAGGING_DB_ENGINE) as session:
+        tag_embedding = TagEmbedding(
+            id=None,
+            category=category,
+            image_url=image_url,
+            image_embeddings=image_embeddings,
+            coordinates=f"POINT({coordinates.lon} {coordinates.lat})",
+            tags={tag.key: tag.value for tag in tags},
+        )
+        session.add(tag_embedding)
+        session.commit()
+
+
 def download_image_s3(image_s3_url: str) -> BytesIO:
     # Parse S3 URL
     parsed_url = urlparse(image_s3_url)