app.py

import gradio as gr
import onnx
import numpy as np
import onnxruntime as ort
from PIL import Image
import cv2

def labels_to_dict(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()
        dictionary = {int(line.split()[0]): line.split()[1] for line in lines}
    return dictionary

labels_dict = labels_to_dict('labels.txt')

model = onnx.load("model.onnx")
session = ort.InferenceSession(model.SerializeToString())

def get_image(path):
    with Image.open(path) as img:
        img = np.array(img.convert('RGB'))
    return img

def preprocess(img):
    img = img / 255.
    img = cv2.resize(img, (256, 256))
    h, w = img.shape[0], img.shape[1]
    y0 = (h - 224) // 2
    x0 = (w - 224) // 2
    img = img[y0 : y0+224, x0 : x0+224, :]  # Corrected cropping
    img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
    img = np.transpose(img, axes=[2, 0, 1])
    img = img.astype(np.float32)
    img = np.expand_dims(img, axis=0)
    return img

def predict(path):
    img = get_image(path)
    img = preprocess(img)
    ort_inputs = {session.get_inputs()[0].name: img}
    preds = session.run(None, ort_inputs)[0]
    preds = np.squeeze(preds)
    a = np.argsort(preds)[::-1]

    results = {labels_dict[a[0]]: preds[a[0]],
              labels_dict[a[1]]: preds[a[1]],
              labels_dict[a[2]]: preds[a[2]]}
        
    # Convert the processed image back to a PIL Image
    processed_img = np.transpose(img[0], axes=[1, 2, 0])
    processed_img = (processed_img * [0.229, 0.224, 0.225]) + [0.485, 0.456, 0.406]
    processed_img = (processed_img * 255).astype(np.uint8)
    processed_img = Image.fromarray(processed_img)
    
    return results, processed_img

iface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="filepath"),
    outputs=[gr.Label(), gr.Image()],    
)
iface.launch()