import gradio as gr
from ultralytics import YOLO
from PIL import Image
import timm
from cods.classif.data import ClassificationDataset
from cods.classif.models import ClassificationModel
from cods.classif.cp import ClassificationConformalizer


def classif(img):
    model_name = "resnet34"
    pretrained_resnet_34 = timm.create_model(model_name, pretrained=True)
    classifier = ClassificationModel(model=pretrained_resnet_34, model_name=model_name)

    val_dataset = ClassificationDataset(...)  # path to imagenet validation set

    val_preds = classifier.build_predictions(
        val_dataset,
        dataset_name="imagenet",
        split_name="cal",
        batch_size=512,
        shuffle=False,
    )

    cc = ClassificationConformalizer(method="lac", preprocess="softmax")
    cc.lbd = 0.9
    conf_cls = cc.conformalize(val_preds)
    return str(conf_cls)


# Load a pretrained YOLOv8n model
model = YOLO("yolov8n.pt")


def main_function(lbd, img):
    results = model(img)  # predict on an image
    r = results[0]
    im_bgr = r.plot()  # BGR-order numpy array
    im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
    new_img = im_rgb
    # res = results[0].save(filename="output.jpg")  # save the image
    # # load image
    # new_img = Image.open("output.jpg")
    return new_img


iface = gr.Interface(
    fn=classif,  # main_function,
    inputs=gr.Image(type="pil"),  # ["slider", gr.Image(type="pil")],
    outputs=gr.Textbox(),  # Image(type="pil"),
    examples=[
        "bus.jpg",  # [0, "bus.jpg"],
    ],
)
iface.launch()