implement logging for response time

main.py

@@ -1,12 +1,24 @@
+import time
 from PIL import Image
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
+from starlette.middleware.base import BaseHTTPMiddleware
 from pydantic import BaseModel, Field
 from transformers import MBartForConditionalGeneration, MBartTokenizerFast
 import torch
 import requests, json, base64
 from io import BytesIO
-from typing import List, Optional
-
+from typing import List, Optional, Tuple
+import os
+import warnings
+import logging
+
+# Suppress specific FutureWarning from huggingface_hub
+warnings.filterwarnings(
+    "ignore",
+    category=FutureWarning,
+    module="huggingface_hub.file_download"
+)
 
 # Initialize FastAPI app with versioning
 app = FastAPI(
@@ -15,18 +27,69 @@ app = FastAPI(
     description="An API for converting text to pictograms, supporting English, French, and Tamil.",
 )
 
+# Set up custom logging format
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    level=logging.INFO
+)
+
+logger = logging.getLogger(__name__)
+
+# Custom Middleware for request time logging
+class RequestTimingMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        start_time = time.time()
+        logger.info(f"Request received: {request.method} {request.url}")
+        
+        # Process the request
+        response = await call_next(request)
+        
+        end_time = time.time()
+        processing_time = end_time - start_time
+        
+        # Log the timing
+        logger.info(f"Request processed: {request.method} {request.url}")
+        logger.info(f"Processing time: {processing_time:.4f}s")
+        
+        # Return the response
+        return response
+
+# Add the middleware to the app
+app.add_middleware(RequestTimingMiddleware)
+
+
+COLORS = {
+    "white": (255, 255, 255),
+    "black": (0, 0, 0),
+    "red": (255, 0, 0),
+    "green": (0, 255, 0),
+    "blue": (0, 0, 255),
+    "yellow": (255, 255, 0),
+    "cyan": (0, 255, 255),
+    "magenta": (255, 0, 255),
+    "gray": (128, 128, 128),
+    "orange": (255, 165, 0),
+    "purple": (128, 0, 128),
+    "brown": (165, 42, 42),
+    "pink": (255, 192, 203),
+    "lime": (0, 255, 0),
+    "teal": (0, 128, 128),
+    "navy": (0, 0, 128)
+    # Add more colors as needed
+}
 
 # Define schemas for requests and responses
+
 class TranslationRequest(BaseModel):
     src: str = Field(..., description="Source text to be translated.")
     language: str = Field(..., description="Language of the source text. Accepted values: 'en', 'fr', 'ta'.")
 
 class TranslationResponse(BaseModel):
-    language: str
-    src: str
-    tgt: Optional[str] = None
-    pictogram_ids: List[Optional[int]] = None
-    image_base64: Optional[str] = None
+    language: str = Field(..., description="Language of the source text.")
+    src: str = Field(..., description="Source text in the original language.")
+    tgt: Optional[str] = Field(None, description="Translated text in the original language.")
+    pictogram_ids: Optional[List[Optional[int]]] = Field(None, description="List of pictogram IDs corresponding to the translation.")
+    image_base64: Optional[str] = Field(None, description="Base64-encoded image of the pictograms, if generated.")
 
 
 # Load the model and tokenizer
@@ -37,20 +100,32 @@ model = MBartForConditionalGeneration.from_pretrained(model_path)
 tokenizer = MBartTokenizerFast.from_pretrained(model_path)
 model = model.to(device)
 
+# Folder to store local pictogram images
+pictogram_folder = 'pictogram_images'
+
 # Load the pictogram dictionary from the JSON file
 with open('pictogram_vocab.json', 'r') as f:
     pictogram_dict = json.load(f)
 
-# Fetch a pictogram image from the ARASAAC API
+# Function to fetch a pictogram image from the ARASAAC API, or use local file if present
 def fetch_pictogram(picto_id: int):
+    # Check if the pictogram image exists locally
+    image_path = os.path.join(pictogram_folder, f"{picto_id}.png")
+    if os.path.exists(image_path):
+        return Image.open(image_path)
+
+    # If the image is not local, request from ARASAAC API
     url = f"https://api.arasaac.org/v1/pictograms/{picto_id}"
     response = requests.get(url)
     if response.status_code == 200:
-        return Image.open(BytesIO(response.content))
+        # Save the image locally
+        img = Image.open(BytesIO(response.content))
+        img.save(image_path)  # Save the image to the local folder
+        return img
     return None
 
-# Generate an image from a sequence of pictogram IDs
-def create_pictogram_image(pictogram_ids):
+# Generate an image from a sequence of pictogram IDs with a customizable background color
+def create_pictogram_image(pictogram_ids, background_color=(255, 0, 0)):  # Default to red background
     pictogram_images = []
     for picto_id in pictogram_ids:
         picto_image = fetch_pictogram(picto_id)
@@ -62,26 +137,56 @@ def create_pictogram_image(pictogram_ids):
     total_width = sum(widths)
     max_height = max(heights)
 
-    final_image = Image.new('RGB', (total_width, max_height))
+    # Create a new image with the specified background color (RGBA with alpha channel)
+    final_image = Image.new('RGBA', (total_width, max_height), background_color + (255,))  # RGBA with alpha channel
     x_offset = 0
     for img in pictogram_images:
-        final_image.paste(img, (x_offset, 0))
+        img = img.convert("RGBA")  # Ensure the image has an alpha channel
+        img_with_bg = Image.new("RGBA", img.size, background_color + (255,))  # Create a solid background
+        img_with_bg.paste(img, (0, 0), img)  # Paste the image on top of the background
+
+        final_image.paste(img_with_bg, (x_offset, 0), img_with_bg)
         x_offset += img.size[0]
 
     return final_image
 
 
 # Endpoint for health check
+@app.get("/")
 @app.get("/health", summary="Health Check", response_description="Health status")
-def health_check():
+async def health_check():
     return {"status": "healthy", "message": "API is up and running"}
 
 # Main translation endpoint
-@app.post("/v1/translate", summary="Translate Text to Pictograms", response_model=TranslationResponse)
-def translate(request: TranslationRequest):
+@app.post("/v1/translate", summary="Translate Text to Pictograms", description="Translates text from a source language to a target language and converts the translation into pictograms. Optionally customize the background color of the generated pictogram images.")
+async def translate(
+        request: TranslationRequest, 
+        backgroundColor: Optional[str] = Query("white"), 
+        backgroundColorRGB: Optional[Tuple[int, int, int]] = Query(None)
+):
+    """
+    Translate the provided source text into pictograms and return a corresponding image.
+    
+    - **src**: The source text to be translated.
+    - **language**: The source language. Accepted values are 'en', 'fr', 'ta'.
+    - **backgroundColor**: (Optional) Background color for the pictogram image, specified by name (e.g., 'red'). Default is 'white'.
+    - **backgroundColorRGB**: (Optional) Background color for the image in RGB format (e.g., (255, 0, 0) for red). This overrides backgroundColor.
+    """
+    
+    # Ensure that both backgroundColor and backgroundColorRGB are not provided simultaneously
+    if backgroundColor and backgroundColorRGB:
+        raise HTTPException(status_code=400, detail="You cannot provide both backgroundColor and backgroundColorRGB at the same time.")
+
+    # Ensure that the given language is a valid one
     if request.language not in ["en", "fr", "ta"]:
         raise HTTPException(status_code=400, detail="Invalid language. Accepted values: 'en', 'fr', 'ta'.")
 
+    # Default to white RGB (255, 255, 255) if neither backgroundColor nor backgroundColorRGB is provided
+    if backgroundColorRGB:  background_color = backgroundColorRGB
+    elif backgroundColor:   background_color = COLORS.get(backgroundColor.lower(), (255, 255, 255))
+    else:                   background_color = (255, 255, 255)  
+
+    # Temporary fix
     # Placeholder for unsupported languages
     if request.language in ["en", "ta"]:
         return TranslationResponse(
@@ -104,7 +209,7 @@ def translate(request: TranslationRequest):
 
     if pictogram_ids:
         # Generate pictogram image
-        final_image = create_pictogram_image(pictogram_ids)
+        final_image = create_pictogram_image(pictogram_ids, background_color)
         if final_image:
             img_byte_arr = BytesIO()
             final_image.save(img_byte_arr, format="PNG")