router/image_clf.py

import os
import time
from typing import Union
from fastapi import APIRouter
from scripts.data_model import ImageInput, ImageOutput
from utils.pipeline import load_model

router = APIRouter()

MODEL_PATH = "ml_models/vit-human-pose-classification/"

@router.post(
    "/image_classification",
    response_model=ImageOutput,
    summary="Image Classification",
    description="Classify the image using a pre-trained model."
)
def image_classification(input: ImageInput)-> ImageOutput:
    try:
        pipe = load_model(MODEL_PATH, is_image_model=True)
        start = time.time()
        output = pipe(input.url)
        end = time.time()
        prediction_time = int((end-start)*1000)

        labels_and_scores = [{"label": x['label'], "score": x['score']} for x in output]

        return ImageOutput(
            user_id=input.user_id,
            url=input.url,
            model_name="vit-human-pose-classification",
            labels=[x['label'] for x in labels_and_scores],
            scores=[x['score'] for x in labels_and_scores],
            prediction_time=prediction_time
        )
   
    except Exception as e:
        return {"error": f"Failed to process text classification: {str(e)}"}, 500