from fastapi import FastAPI, Response, HTTPException
from datasets import load_dataset
from io import BytesIO
from PIL import Image

app = FastAPI()

dataset = load_dataset("visionLMsftw/vibe-testing-samples", split="train")
id_to_image = {example["ex_id"]: example["image"] for example in dataset}

print(id_to_image)

@app.get("/image/{image_id}")
def get_image(image_id: int):
    if image_id not in id_to_image:
        raise HTTPException(status_code=404, detail="Image not found")

    image: Image.Image = id_to_image[image_id].convert("RGB")
    buffer = BytesIO()
    image.save(buffer, format="JPEG", quality=85)
    buffer.seek(0)
    return Response(content=buffer.read(), media_type="image/jpeg")

@app.get("/ids")
def list_ids():
    return list(id_to_image.keys())