CODEWITHSUNDEEP
pythonfastapi
pavel hrablis
pavel hrablis

Reputation: 141

How to return image and json in one response in fastapi?

I get an image, change it, then it is classified using a neural network, should return a new image and json with a response. How to do it with one endpoint? image is returned with Streaming Response but how to add json to it?

import io
from starlette.responses import StreamingResponse

app = FastAPI()

@app.post("/predict")
def predict(file: UploadFile = File(...)):
    img = file.read()
    new_image = prepare_image(img)
    result = predict(new_image)
    return StreamingResponse(io.BytesIO(new_image.tobytes()), media_type="image/png")

Upvotes: 8

Views: 11837

Answers (3)

Jees Antony
Jees Antony

Reputation: 3

You can use this code to return an Image and json altogether when using yolov8 model, hope this work with other models too, please modify the image generation part

@app.post("/detect")
async def detect_and_return_image(image_file: UploadFile = File(...)):
    """
    Handler of /detect POST endpoint
    Receives uploaded file with a name "image_file",
    passes it through YOLOv8 object detection
    network and returns an array of bounding boxes.
    :return: a JSON array of objects bounding
    boxes in format
    [[x1,y1,x2,y2,object_type,probability],..]
    """
    buf = await image_file.read()
    boxes, class_prob = detect_objects_on_image(Image.open(BytesIO(buf)))
    print(f'class proba {class_prob}')
    annotated_image = annotate_image(Image.open(BytesIO(buf)), boxes)
    return {
        "annotated_image": image_to_base64(annotated_image),
        "class_prob": class_prob
    }


def detect_objects_on_image(image):
    """
    Function receives an image,
    passes it through YOLOv8 neural network
    and returns an array of detected objects
    and their bounding boxes
    :param image: Input image
    :return: Array of bounding boxes in format
    [[x1,y1,x2,y2,object_type,probability],..]
    """
    model = YOLO('best.pt')
    results = model.predict(image)
    result = results[0]
    output = []
    class_prob = []
    for box in result.boxes:
        x1, y1, x2, y2 = [round(x) for x in box.xyxy[0].tolist()]
        class_id = box.cls[0].item()
        prob = round(box.conf[0].item(), 2)
        output.append([x1, y1, x2, y2, result.names[class_id], prob])
        class_prob.append([result.names[class_id], prob])
    return output, class_prob


def annotate_image(image, boxes):
    """
    Function annotates the image with bounding boxes.
    :param image: Input image
    :param boxes: Array of bounding boxes in format
    [[x1,y1,x2,y2,object_type,probability],..]
    :return: Annotated image
    """
    # Draw bounding boxes on the image
    draw = ImageDraw.Draw(image)
    for box in boxes:
        x1, y1, x2, y2, object_type, probability = box
        draw.rectangle([(x1, y1), (x2, y2)], outline="red", width=3)
        draw.text((x1, y1 - 10), f"{object_type} ({probability})", fill="red")

    return image


def save_annotated_image(image):
    """
    Function saves the annotated image and returns
    the image as a response.
    :param image: Annotated image
    :return: StreamingResponse with the image
    """
    output_buffer = BytesIO()
    image.save(output_buffer, format="PNG")
    output_buffer.seek(0)
    return StreamingResponse(output_buffer, media_type="image/png")

def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

To get full code please visit github here

Upvotes: 0

PankajSanwal
PankajSanwal

Reputation: 1019

I was having the same issue, although, my file was stored locally but still I have to return JSON, and Image in a single response.

This worked for me, much neater and shorter:

@app.post("/ImgAndJSON")
# Postmsg is a Pydantic model having 1 str field
def ImgAndJSON(message:PostMsg):

    results={"message":"This is just test message"}

    return FileResponse('path/to/file.png',headers=results)

Upvotes: 3

pavel hrablis
pavel hrablis

Reputation: 141

I added json to response headers. change from:

@app.post("/predict")
def predict(file: UploadFile = File(...)):
    img = file.read()
    new_image = prepare_image(img)
    result = predict(new_image)
    return StreamingResponse(io.BytesIO(new_image.tobytes()), media_type="image/png")

to

@app.post("/predict/")
def predict(file: UploadFile = File(...)):
    file_bytes = file.file.read()
    image = Image.open(io.BytesIO(file_bytes))
    new_image = prepare_image(image)
    result = predict(image)
    bytes_image = io.BytesIO()
    new_image.save(bytes_image, format='PNG')
    return Response(content = bytes_image.getvalue(), headers = result, media_type="image/png")

Upvotes: 6

Related Questions