src/apis/img_processing_api.py

"""
Created By: ishwor subedi
Date: 2024-08-13
"""
from PIL import Image
from fastapi import HTTPException, UploadFile, File
from fastapi.routing import APIRouter
from six import BytesIO

from src.models.models import ImageCaptionRequest, ImageGenerationRequest, LanguageTranslationRequest
from src.pipeline.image_processing_pipeline import ImageProcessingPipeline

image_processing_pipeline = ImageProcessingPipeline()
image_processing_router = APIRouter()


@image_processing_router.post("/generate_image")
async def generate_image(request: ImageGenerationRequest):
    try:
        image = image_processing_pipeline.generate_image(
            request.prompt,
            request.negative_prompt,
            request.style,
            request.use_negative_prompt,
            request.num_inference_steps,
            request.num_images_per_prompt,
            request.seed,
            request.width,
            request.height,
            request.guidance_scale,
            request.randomize_seed
        )
        base_64_image = BytesIO()
        image.save(base_64_image, format='PNG')
        base_64_image = base_64_image.getvalue()
        json = {'image': base_64_image, 'status_code': 200}

        return json
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@image_processing_router.post("/generate_caption")
async def generate_caption(request: ImageCaptionRequest, image: UploadFile = File(...)):
    image = Image.open(BytesIO(await image.read())).convert('RGB')

    try:
        caption = image_processing_pipeline.generate_caption(
            image,
            request.prompt,
            request.temperature,
            request.length_penalty,
            request.repetition_penalty,
            request.max_length,
            request.min_length,
            request.top_p
        )
        json = {'caption': caption, 'status_code': 200}
        return json
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))