frist commit

router/disaster.py

@@ -1,6 +1,5 @@
 import os
 import time
-from typing import Union
 from fastapi import APIRouter
 from scripts.data_model import ClassificationInput, ClassificationOutput
 from utils.pipeline import load_model
@@ -17,6 +16,16 @@ MODEL_PATH = os.path.join(BASE_DIR, "ml-models", "tinybert-disaster-tweet/")
     description="Classify the disaster of a given text using a pre-trained model."
 )
 def disaster_classification(input: ClassificationInput)-> ClassificationOutput:
+    """
+    Classify the disaster of a given text using a pre-trained model.
+
+    Args:
+        input (ClassificationInput): The input data containing the user_id and text.
+
+    Returns:
+        ClassificationOutput: The output data containing the user_id, text, model_name, sentiment, score, and prediction_time.
+
+    """
     try:
         pipe = load_model(MODEL_PATH)
         start = time.time()

router/image_clf.py

@@ -1,13 +1,8 @@
 import os
 import time
-import requests
-from PIL import Image
-from io import BytesIO
-
 from fastapi import APIRouter, HTTPException
 from scripts.data_model import ImageInput, ImageOutput
 from utils.pipeline import load_model
-from utils.log import logger
 
 router = APIRouter()
 
@@ -21,20 +16,27 @@ MODEL_PATH = os.path.join(BASE_DIR, "ml-models", "vit-human-pose-classification/
     description="Classify the image using a pre-trained model."
 )
 def image_classification(input: ImageInput)-> ImageOutput:
+    """
+    Classify the image using a pre-trained model.
+    
+    Args:
+        input (ImageInput): The input data containing the image URL and user ID.
+    
+    Returns:
+        ImageOutput: The output data containing the labels, scores, prediction time, and other info.
+    """
     try:
         pipe = load_model(MODEL_PATH, is_image_model=True)
-        image = input.url
-        
-        logger.info(f"Image URLs: {image[-1]}")
+        urls = [str(x) for x in input.url]
         
         start = time.time()
-        output = pipe(image[-1])
+        output = pipe(urls)
         end = time.time()
         
         prediction_time = int((end-start)*1000)
         
-        labels = [x["label"] for x in output]
-        scores = [x["score"] for x in output]
+        labels = [x[0]['label'] for x in output]
+        scores = [x[0]['score'] for x in output]
 
         return ImageOutput(
             user_id=input.user_id,
@@ -46,4 +48,4 @@ def image_classification(input: ImageInput)-> ImageOutput:
         )
    
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to process image classification: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to process image classification: {e}")

router/sentiment.py

@@ -1,6 +1,5 @@
 import os
 import time
-from typing import Union
 from fastapi import APIRouter
 from scripts.data_model import ClassificationInput, ClassificationOutput
 from utils.pipeline import load_model
@@ -17,6 +16,16 @@ MODEL_PATH = os.path.join(BASE_DIR, "ml-models", "tinybert-sentiment-analysis/")
     description="Classify the sentiment of a given text using a pre-trained model."
 )
 def sentiment_classification(input: ClassificationInput)-> ClassificationOutput:
+    """
+    Classify the sentiment of a given text using a pre-trained model.
+
+    Args:
+        input (ClassificationInput): The input data containing the user_id and text.
+
+    Returns:
+        ClassificationOutput: The output data containing the user_id, text, model_name, sentiment, score, and prediction_time.
+
+    """
     try:
         pipe = load_model(MODEL_PATH)
         start = time.time()

router/upload_image.py

@@ -1,8 +0,0 @@
-from fastapi import APIRouter
-from fastapi.responses import FileResponse
-
-# router = APIRouter()
-
-# @router.get("/image/{image_name}")
-# async def get_image(image_name: str):
-#     return FileResponse(f"image/{image_name}")

scripts/data_model.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, HttpUrl
 from typing import List
 
 
@@ -8,7 +8,7 @@ class ClassificationInput(BaseModel):
 
 class ImageInput(BaseModel):
     user_id: str
-    url: List[str]
+    url: List[HttpUrl]
 
 
 class ClassificationOutput(BaseModel):
@@ -21,7 +21,7 @@ class ClassificationOutput(BaseModel):
 
 class ImageOutput(BaseModel):
     user_id: str
-    url: List[str]
+    url: List[HttpUrl]
     model_name: str
     label: List[str]
     score: List[float]

scripts/s3.py

@@ -14,6 +14,16 @@ s3 = boto3.client(
 )
 
 def download_model_from_s3(local_path: str, s3_prefix: str):
+    """
+    Downloads a model from S3 to the specified local path.
+
+    Args:
+        local_path (str): The local path to download the model to.
+        s3_prefix (str): The S3 prefix of the model to download.
+
+    Raises:
+        RuntimeError: If there is an error downloading the model from S3.
+    """
     try:
         if os.path.exists(local_path) and os.listdir(local_path):
             logger.info(f"Model {local_path} already exists. Skipping download.")
@@ -41,6 +51,18 @@ def upload_image_to_s3(
         s3_prefix="ml-images", 
         object_name=None
     ):
+    """
+    Uploads an image to S3 and returns a presigned URL for the object.
+
+    Args:
+        file_name (str): The file name of the image to upload.
+        s3_prefix (str): The S3 prefix to use for the object name.
+        object_name (str, optional): The object name to use for the S3 key.
+            If not provided, the object name will be the same as the file name.
+
+    Returns:
+        str: The presigned URL for the S3 object.
+    """
     if object_name is None:
         object_name = os.path.basename(file_name)
 

utils/log.py

@@ -1,6 +1,15 @@
 import logging
 
 def get_logger(name: str):
+    """
+    Gets a logger with the given name, and configures it if it has never been configured before.
+
+    Args:
+        name (str): The name of the logger to get.
+
+    Returns:
+        A configured logger with the given name.
+    """
     logger = logging.getLogger(name)
     if not logger.handlers: 
         handler = logging.StreamHandler()