Spaces:
Sleeping
Sleeping
File size: 1,607 Bytes
db2db2a 665e8bb db2db2a 7312439 db2db2a c034083 db2db2a c034083 665e8bb db2db2a c034083 db2db2a 665e8bb db2db2a aa385b3 c034083 db2db2a aa385b3 db2db2a c034083 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import os
import time
from fastapi import APIRouter, HTTPException
from scripts.data_model import ImageInput, ImageOutput
from utils.pipeline import load_model
router = APIRouter()
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODEL_PATH = os.path.join(BASE_DIR, "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:
"""
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)
urls = [str(x) for x in input.url]
start = time.time()
output = pipe(urls)
end = time.time()
prediction_time = int((end-start)*1000)
labels = [x[0]['label'] for x in output]
scores = [x[0]['score'] for x in output]
return ImageOutput(
user_id=input.user_id,
url=input.url,
model_name="vit-human-pose-classification",
label=labels,
score=scores,
prediction_time=prediction_time
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to process image classification: {e}")
|