Spaces:
Runtime error
Runtime error
File size: 1,431 Bytes
d49f9c6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
from fastapi import APIRouter
from datasets import load_dataset
from ast import literal_eval
from pydantic import BaseModel
from typing import Dict
from io import BytesIO
from PIL import Image
import base64
from config import settings
from huggingface_hub import login
router = APIRouter()
login(settings.huggingface_key)
class ImageResponse(BaseModel):
image_data: str
ground_truth_data: Dict
def encode_pil_image(image: Image) -> str:
buffer = BytesIO()
image.save(buffer, format='JPEG')
img_data = buffer.getvalue()
return base64.b64encode(img_data).decode('utf-8')
@router.get("/dataset_info")
async def get_dataset_info():
dataset = load_dataset(settings.dataset_name)
splits = []
for split in dataset.keys():
split = {
"name": split,
"number_of_rows": len(dataset[split])
}
splits.append(split)
result = {
"dataset": settings.dataset_name,
"splits": splits
}
return result
@router.get("/ground_truth", response_model=ImageResponse)
async def get_ground_truth() -> ImageResponse:
dataset = load_dataset(settings.dataset_name)
example = dataset['test'][0]
image = example['image']
encoded_img = encode_pil_image(image)
ground_truth = example['ground_truth']
data = literal_eval(ground_truth)['gt_parse']
return ImageResponse(image_data=encoded_img, ground_truth_data=data) |