Convert COCO images to KSserve payloads

coco_to_kserve.py

import os
import json
import numpy as np
from PIL import Image

def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result

    
INPATH = "test2017/" # set this to whatever path contains your raw COCO jpgs
OUTPATH= "test2017_proc/" # set this to whatever output directory you want to save the KServe json files to

for f in os.listdir(INPATH):
    # load image
    im = Image.open(INPATH+f)
    
    # pad to 640, 640 square
    im = expand2square(im, (0,0,0))
    
    # convert to np array of correct shape
    arr = np.transpose(np.array(im), (2, 0, 1))
    arr = np.expand_dims(arr, axis=0)
    
    # write to json
    row = {"name": "images", "shape":arr.shape, "datatype":"FP32"}
    row["data"] = arr.tolist()
    datajson = {"inputs":[row]}
    
    outname = f.replace(".jpg",".json")
    with open(OUTPATH+outname, "w") as outfile:
        json.dump(datajson, outfile)
    
    # remove this break to process all images
    break