src/modules/image_analysis.py

import json
from typing import Literal

from fireworks.llm import LLM
from configs.load_config import PROMPT_LIBRARY, APP_STEPS_CONFIGS
from pydantic import BaseModel
import io
import base64


class IncidentAnalysis(BaseModel):
    description: str
    location: Literal["front-left", "front-right", "back-left", "back-right"]
    severity: Literal["minor", "moderate", "major"]
    license_plate: str


def get_llm(api_key: str, model: str, temperature: float) -> LLM:
    return LLM(
        model=model,
        temperature=temperature,
        deployment_type="serverless",
        api_key=api_key,
    )


def pil_to_base64_dict(pil_image):
    """Convert PIL image to the format expected by analyze_damage_image"""
    if pil_image is None:
        return None

    buffered = io.BytesIO()
    if pil_image.mode != "RGB":
        pil_image = pil_image.convert("RGB")

    pil_image.save(buffered, format="JPEG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

    return {"image": pil_image, "path": "uploaded_image.jpg", "base64": img_base64}


def analyze_damage_image(image, api_key: str, prompt: str = "advanced"):
    """
    Analyze the damage in an image using the Fireworks VLM model.
    """
    assert (
        prompt in PROMPT_LIBRARY["vision_damage_analysis"].keys()
    ), f"Invalid prompt choose from {list(PROMPT_LIBRARY['vision_damage_analysis'].keys())}"

    prompt_text = PROMPT_LIBRARY["vision_damage_analysis"][prompt]

    llm = get_llm(
        api_key=api_key,
        model=APP_STEPS_CONFIGS.analyze_damage_image.model,
        temperature=APP_STEPS_CONFIGS.analyze_damage_image.temperature,
    )
    response = llm.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{image['base64']}"
                        },
                    },
                    {"type": "text", "text": prompt_text},
                ],
            }
        ],
        response_format={
            "type": "json_object",
            "schema": IncidentAnalysis.model_json_schema(),
        },
    )

    result = json.loads(response.choices[0].message.content)
    return result