Spaces:
Running
Running
| from transformers import pipeline | |
| from pydantic import BaseModel | |
| import logging | |
| from fastapi import Request, HTTPException | |
| import json | |
| from typing import Optional | |
| class ImageClassificationRequest(BaseModel): | |
| inputs: str | |
| parameters: Optional[dict] = None | |
| class ImageClassificationTaskService: | |
| __logger: logging.Logger | |
| __task_name: str | |
| def __init__(self, logger: logging.Logger, task_name: str = "image-classification"): | |
| self.__logger = logger | |
| self.__task_name = task_name | |
| async def get_image_classification_request( | |
| self, | |
| request: Request | |
| ) -> ImageClassificationRequest: | |
| content_type = request.headers.get("content-type", "") | |
| if content_type.startswith("application/json"): | |
| data = await request.json() | |
| return ImageClassificationRequest(**data) | |
| if content_type.startswith("application/x-www-form-urlencoded"): | |
| raw = await request.body() | |
| try: | |
| data = json.loads(raw) | |
| return ImageClassificationRequest(**data) | |
| except Exception: | |
| try: | |
| data = json.loads(raw.decode("utf-8")) | |
| return ImageClassificationRequest(**data) | |
| except Exception: | |
| raise HTTPException(status_code=400, detail="Invalid request body") | |
| raise HTTPException(status_code=400, detail="Unsupported content type") | |
| async def classify( | |
| self, | |
| request: Request, | |
| model_name: str | |
| ): | |
| imageRequest: ImageClassificationRequest = await self.get_image_classification_request(request) | |
| try: | |
| pipe = pipeline(self.__task_name, model=model_name) | |
| except Exception as e: | |
| self.__logger.error(f"Failed to load model '{model_name}': {str(e)}") | |
| raise HTTPException( | |
| status_code=404, | |
| detail=f"Model '{model_name}' could not be loaded: {str(e)}" | |
| ) | |
| try: | |
| if self.__task_name == "zero-shot-image-classification": | |
| candidate_labels = [] | |
| if imageRequest.parameters: | |
| candidate_labels = imageRequest.parameters.get('candidate_labels', []) | |
| if isinstance(candidate_labels, str): | |
| candidate_labels = [label.strip() for label in candidate_labels.split(',')] | |
| result = pipe(imageRequest.inputs, candidate_labels=candidate_labels) | |
| else: # image classification | |
| result = pipe(imageRequest.inputs) | |
| except Exception as e: | |
| self.__logger.error(f"Inference failed for model '{model_name}': {str(e)}") | |
| raise HTTPException( | |
| status_code=500, | |
| detail=f"Inference failed: {str(e)}" | |
| ) | |
| return result |