additional tasks

src/classification.py

@@ -0,0 +1,82 @@
+from transformers import pipeline
+from pydantic import BaseModel
+import logging
+from fastapi import Request, HTTPException
+import json
+from typing import Optional
+
+
+class ClassificationRequest(BaseModel):
+    inputs: str
+    parameters: Optional[dict] = None
+
+class ClassificationTaskService:    
+
+    __logger: logging.Logger
+    __task_name: str
+
+    def __init__(self, logger: logging.Logger, task_name: str):
+        self.__logger = logger
+        self.__task_name = task_name
+
+    async def get_classification_request(
+        self,
+        request: Request
+    )  -> ClassificationRequest:
+        content_type = request.headers.get("content-type", "")
+        if content_type.startswith("application/json"):
+            data = await request.json()
+            return ClassificationRequest(**data)
+        if content_type.startswith("application/x-www-form-urlencoded"):
+            raw = await request.body()
+            try:
+                data = json.loads(raw)
+                return ClassificationRequest(**data)
+            except Exception:
+                try:
+                    data = json.loads(raw.decode("utf-8"))
+                    return ClassificationRequest(**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
+    ): 
+
+        classificationRequest: ClassificationRequest = await self.get_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" or self.__task_name == "zero-shot-classification":
+                candidate_labels = []
+
+                if classificationRequest.parameters:
+                    candidate_labels = classificationRequest.parameters.get('candidate_labels', [])
+                    if isinstance(candidate_labels, str):
+                        candidate_labels = [label.strip() for label in candidate_labels.split(',')]
+                result = pipe(classificationRequest.inputs, candidate_labels=candidate_labels)
+
+            else:  # pretrained classification
+                result = pipe(classificationRequest.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        

src/image_classification.py

@@ -0,0 +1,82 @@
+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        

src/main.py

@@ -8,21 +8,15 @@
 #  @license    Pimcore Open Core License (POCL)
 # -------------------------------------------------------------------
 
-import os
 import torch
 
-from fastapi import FastAPI, Path, Depends, HTTPException, UploadFile, Form, File, status, Request
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from pydantic import BaseModel
-from typing import Annotated
-import json
-
+from fastapi import FastAPI, Path, Request
 import logging
 import sys
-import base64
-
 
-from transformers import pipeline
+from .translation_task import TranslationTaskService
+from .classification import ClassificationTaskService
+from .text_to_image import TextToImageTaskService
 
 app = FastAPI(
     title="Pimcore Local Inference Service",
@@ -51,14 +45,6 @@ class StreamToLogger(object):
 sys.stdout = StreamToLogger(logger, logging.INFO)
 sys.stderr = StreamToLogger(logger, logging.ERROR)
 
-
-
-class ResponseModel(BaseModel):
-    """ Default response model for endpoints. """
-    message: str
-    success: bool = True
-
-
 @app.get("/gpu_check")
 async def gpu_check():
     """ Check if a GPU is available """
@@ -73,41 +59,9 @@ async def gpu_check():
     return {'success': True, 'gpu': gpu}
 
 
-from typing import Optional
-
-
-
 # =========================
 # Translation Task
 # =========================
-
-class TranslationRequest(BaseModel):
-    inputs: str
-    parameters: Optional[dict] = None
-    options: Optional[dict] = None
-
-async def get_translation_request(
-    request: Request
-)  -> TranslationRequest: 
-    content_type = request.headers.get("content-type", "")
-    if content_type.startswith("application/json"):
-        data = await request.json()
-        return TranslationRequest(**data)
-    if content_type.startswith("application/x-www-form-urlencoded"):
-        raw = await request.body()
-        try:
-            data = json.loads(raw)
-            return TranslationRequest(**data)
-        except Exception:
-            try:
-                data = json.loads(raw.decode("utf-8"))
-                return TranslationRequest(**data)
-            except Exception:
-                raise HTTPException(status_code=400, detail="Invalid request body")
-    raise HTTPException(status_code=400, detail="Unsupported content type")
-
-
-
 @app.post(
     "/translation/{model_name:path}/", 
     openapi_extra={
@@ -138,60 +92,13 @@ async def translate(
         list: The translation result(s) as returned by the pipeline.
     """
 
-    translationRequest: TranslationRequest = await get_translation_request(request)
-
-    try:
-        pipe = pipeline("translation", model=model_name)
-    except Exception as e:
-        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:       
-        result = pipe(translationRequest.inputs, **(translationRequest.parameters or {}))
-    except Exception as e:
-        logger.error(f"Inference failed for model '{model_name}': {str(e)}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"Inference failed: {str(e)}"
-        )
-
-    return result
+    translationTaskService = TranslationTaskService(logger)
+    return await translationTaskService.translate(request, model_name)
 
 
 # =========================
 # Zero-Shot Image Classification Task
 # =========================
-
-
-class ZeroShotImageClassificationRequest(BaseModel):
-    inputs: str
-    parameters: Optional[dict] = None
-
-async def get_zero_shot_image_classification_request(
-    request: Request
-)  -> ZeroShotImageClassificationRequest: 
-    content_type = request.headers.get("content-type", "")
-    if content_type.startswith("application/json"):
-        data = await request.json()
-        return ZeroShotImageClassificationRequest(**data)
-    if content_type.startswith("application/x-www-form-urlencoded"):
-        raw = await request.body()
-        try:
-            data = json.loads(raw)
-            return ZeroShotImageClassificationRequest(**data)
-        except Exception:
-            try:
-                data = json.loads(raw.decode("utf-8"))
-                return ZeroShotImageClassificationRequest(**data)
-            except Exception:
-                raise HTTPException(status_code=400, detail="Invalid request body")
-    raise HTTPException(status_code=400, detail="Unsupported content type")
-
-
-
 @app.post(
     "/zero-shot-image-classification/{model_name:path}/",
     openapi_extra={
@@ -222,58 +129,126 @@ async def zero_shot_image_classification(
         list: The classification result(s) as returned by the pipeline.
     """
 
-    zeroShotRequest: ZeroShotImageClassificationRequest = await get_zero_shot_image_classification_request(request)
+    zeroShotTask = ClassificationTaskService(logger, 'zero-shot-image-classification')
+    return await zeroShotTask.classify(request, model_name)
+
 
-    try:
-        pipe = pipeline("zero-shot-image-classification", model=model_name)
-    except Exception as e:
-        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)}"
-        )
+# =========================
+# Image Classification Task
+# =========================
+@app.post(
+    "/image-classification/{model_name:path}/",
+    openapi_extra={
+        "requestBody": {
+            "content": {
+                "application/json": {
+                    "example": {
+                        "inputs": "base64_encoded_image_string"
+                    }
+                }
+            }
+        }        
+    }
+)
+async def image_classification(
+    request: Request,
+    model_name: str = Path(
+        ...,
+        description="The name of the image classification model (e.g., pimcore/car-countries-classification)",
+        example="pimcore/car-countries-classification"
+    )
+    ):
+    """
+    Execute image classification tasks.
 
-    try:       
-        candidate_labels = []
-        if zeroShotRequest.parameters:
-            candidate_labels = zeroShotRequest.parameters.get('candidate_labels', [])
-            if isinstance(candidate_labels, str):
-                candidate_labels = [label.strip() for label in candidate_labels.split(',')]
-        result = pipe(zeroShotRequest.inputs, candidate_labels=candidate_labels)
-    except Exception as e:
-        logger.error(f"Inference failed for model '{model_name}': {str(e)}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"Inference failed: {str(e)}"
-        )
+    Returns:
+        list: The classification result(s) as returned by the pipeline.
+    """
 
-    return result
+    imageTask = ClassificationTaskService(logger, 'image-classification')
+    return await imageTask.classify(request, model_name)
 
 
 
 # =========================
-# Image to Text Task
+# Zero-Shot Text Classification Task
 # =========================
+@app.post(
+    "/zero-shot-text-classification/{model_name:path}/",
+    openapi_extra={
+        "requestBody": {
+            "content": {
+                "application/json": {
+                    "example": {
+                        "inputs": "text to classify",
+                        "parameters": {"candidate_labels": "green, yellow, blue, white, silver"}
+                    }
+                }
+            }
+        }        
+    }
+)
+async def zero_shot_text_classification(
+    request: Request,
+    model_name: str = Path(
+        ...,
+        description="The name of the zero-shot text classification model (e.g., facebook/bart-large-mnli)",
+        example="facebook/bart-large-mnli"
+    )
+    ):
+    """
+    Execute zero-shot text classification tasks.
 
+    Returns:
+        list: The classification result(s) as returned by the pipeline.
+    """
 
-async def get_encoded_image(
-    request: Request
-)  -> str: 
-    content_type = request.headers.get("content-type", "")
-    if content_type.startswith("multipart/form-data"):
-        form = await request.form()
-        image = form.get("image")
-        if image:
-            image_bytes = await image.read()
-            return base64.b64encode(image_bytes).decode("utf-8")
-    if content_type.startswith("image/"):
-        image_bytes = await request.body()
-        return base64.b64encode(image_bytes).decode("utf-8")
+    zeroShotTask = ClassificationTaskService(logger, 'zero-shot-classification')
+    return await zeroShotTask.classify(request, model_name)
 
-    raise HTTPException(status_code=400, detail="Unsupported content type")
 
+# =========================
+# Text Classification Task
+# =========================
+@app.post(
+    "/text-classification/{model_name:path}/",
+    openapi_extra={
+        "requestBody": {
+            "content": {
+                "application/json": {
+                    "example": {
+                        "inputs": "text to classify"
+                    }
+                }
+            }
+        }        
+    }
+)
+async def text_classification(
+    request: Request,
+    model_name: str = Path(
+        ...,
+        description="The name of the text classification model (e.g., pimcore/car-class-classification)",
+        example="pimcore/car-class-classification"
+    )
+    ):
+    """
+    Execute text classification tasks.
+
+    Returns:
+        list: The classification result(s) as returned by the pipeline.
+    """
+
+    textTask = ClassificationTaskService(logger, 'text-classification')
+    return await textTask.classify(request, model_name)
 
 
+
+
+
+# =========================
+# Image to Text Task
+# =========================
 @app.post(
     "/image-to-text/{model_name:path}/",
     openapi_extra={
@@ -311,24 +286,5 @@ async def image_to_text(
         list: The generated text as returned by the pipeline.
     """
 
-    encoded_image = await get_encoded_image(request)
-
-    try:
-        pipe = pipeline("image-to-text", model=model_name, use_fast=True)
-    except Exception as e:
-        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:       
-        result = pipe(encoded_image)
-    except Exception as e:
-        logger.error(f"Inference failed for model '{model_name}': {str(e)}")
-        raise HTTPException(
-            status_code=500,
-            detail=f"Inference failed: {str(e)}"
-        )
-
-    return result
+    imageToTextTask = TextToImageTaskService(logger)
+    return await imageToTextTask.extract(request, model_name)

src/text_to_image.py

@@ -0,0 +1,56 @@
+from transformers import pipeline
+import logging
+from fastapi import Request, HTTPException
+import base64
+
+
+class TextToImageTaskService:    
+
+    __logger: logging.Logger
+
+    def __init__(self, logger: logging.Logger):
+        self.__logger = logger
+
+    async def get_encoded_image(
+        self,
+        request: Request
+    )  -> str: 
+        content_type = request.headers.get("content-type", "")
+        if content_type.startswith("multipart/form-data"):
+            form = await request.form()
+            image = form.get("image")
+            if image:
+                image_bytes = await image.read()
+                return base64.b64encode(image_bytes).decode("utf-8")
+        if content_type.startswith("image/"):
+            image_bytes = await request.body()
+            return base64.b64encode(image_bytes).decode("utf-8")
+
+        raise HTTPException(status_code=400, detail="Unsupported content type")
+
+    async def extract(
+        self,
+        request: Request,
+        model_name: str
+    ): 
+        encoded_image = await self.get_encoded_image(request)
+
+        try:
+            pipe = pipeline("image-to-text", model=model_name, use_fast=True)
+        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:       
+            result = pipe(encoded_image)
+        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        

src/translation_task.py

@@ -0,0 +1,67 @@
+from transformers import pipeline
+from pydantic import BaseModel
+import logging
+from fastapi import Request, HTTPException
+import json
+from typing import Optional
+
+class TranslationRequest(BaseModel):
+    inputs: str
+    parameters: Optional[dict] = None
+    options: Optional[dict] = None
+
+class TranslationTaskService:
+
+    __logger: logging.Logger
+
+    def __init__(self, logger: logging.Logger):
+        self.__logger = logger
+
+    async def get_translation_request(
+        self, 
+        request: Request
+    )  -> TranslationRequest: 
+        content_type = request.headers.get("content-type", "")
+        if content_type.startswith("application/json"):
+            data = await request.json()
+            return TranslationRequest(**data)
+        if content_type.startswith("application/x-www-form-urlencoded"):
+            raw = await request.body()
+            try:
+                data = json.loads(raw)
+                return TranslationRequest(**data)
+            except Exception:
+                try:
+                    data = json.loads(raw.decode("utf-8"))
+                    return TranslationRequest(**data)
+                except Exception:
+                    raise HTTPException(status_code=400, detail="Invalid request body")
+        raise HTTPException(status_code=400, detail="Unsupported content type")
+
+
+    async def translate(
+            self, 
+            request: Request, 
+            model_name: str
+        ):
+
+        translationRequest: TranslationRequest = await self.get_translation_request(request)
+
+        try:
+            pipe = pipeline("translation", 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:       
+            result = pipe(translationRequest.inputs, **(translationRequest.parameters or {}))
+            return result
+        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)}"
+            )