app.py

import gradio as gr
from utils import *
from PIL import Image
from gen_dataset import preds_to_data, data_to_preds
from algo import final_say, confidence_score
import os

seed_everything()

def predict(image):
    
    c1, c2, v1, v2, out1, out2 = get_predictions(image)

    probas = preds_to_data(c1, c2).unsqueeze(0).to(DEVICE)
    prediction = Morpheus(probas).argmax(1).item()
    out3 = list(CLASSES_1.keys())[prediction]
    
    pred = final_say(v1, v2, out1, out2, out3)
    confidence = confidence_score(pred, probas)
    
    output = {
        "result": pred,
        "m1_confidence": v1,
        "m2_confidence": v2,
        "probabilities": probas.tolist(),
        "embedding_confidence": confidence,
    }
    
    print(output)
    
    return output

title = "Junk Judge - A Model to Classify Garbage Images"
description = '<center><img src="https://github.com/maksymalist/maksymalist/assets/79988159/288ca815-8aef-4a67-b4b2-7ab12b637c02" alt="Junk Judge" width="200px" height="200px"></center>'

gradio = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    flagging_options=["blurry", "incorrect", "other"],
    outputs=gr.outputs.JSON(),
    title=title,
    description=description,
    theme=gr.themes.Base(primary_hue="lime"),
    css="#component-1 {justify-content: center;align-items: center;flex-direction: column; width: 100%} #component-5 {justify-content: center;align-items: center;flex-direction: column; width: 100%} #component-6 {width: 100%; max-width: 900px} #component-11 {width: 100%; max-width: 900px}"
    
)
gradio.launch()


# VISUALIZATION

# fig, ax = plt.subplots(1, 3, figsize=(10, 10))

# ax[0].pie(v1.values(), labels=v1.keys())
# ax[0].margins(x=20, y=10)
# ax[0].set_title("First Model")

# ax[1].pie(v2.values(), labels=v2.keys())
# ax[1].margins(x=20, y=10)
# ax[1].set_title("Second Model")

# ax[2].imshow(img1)
# ax[2].axis("off")
# ax[2].margins(x=20, y=10)
# ax[2].set_title(final_verdict)

# plt.show()