import gradio as gr
from utils import Predictor, setup

CAR_PART_SEG_OUT, CAR_DAM_DET_OUT, CAR_SEV_DET_OUT, SAMPLE = setup()

car_part_predictor = Predictor(CAR_PART_SEG_OUT)
car_dam_detector = Predictor(CAR_DAM_DET_OUT)
car_sev_detector = Predictor(CAR_SEV_DET_OUT)

car_part_ui = gr.Interface(
    fn=car_part_predictor.transform,
    inputs=[ 
        gr.components.Image(type="filepath", label="Input Image"),
        gr.components.Checkbox(value=True, label="Annotate Boxes"),
        gr.components.Checkbox(value=False, label="Annotate Masks"),
        ],
    outputs=gr.components.Image(type="numpy", label="Output Image"),
    title="Car Parts Detector",
    examples=SAMPLE["car-parts-seg"],
    cache_examples=False,
)

car_dam_ui = gr.Interface(
    fn=car_dam_detector.transform,
    inputs=[ 
        gr.components.Image(type="filepath", label="Input Image"),
        gr.components.Checkbox(value=True, label="Annotate Boxes"),
        gr.components.Checkbox(value=False, label="Annotate Masks"),
        ],
    outputs=gr.components.Image(type="numpy", label="Output Image"),
    title="Car Damage Detector",
    examples=SAMPLE["car-dam-det"],
    cache_examples=False,
)

car_sev_ui = gr.Interface(
    fn=car_sev_detector.transform,
    inputs=[ 
        gr.components.Image(type="filepath", label="Input Image"),
        gr.components.Checkbox(value=True, label="Annotate Boxes"),
        gr.components.Checkbox(value=False, label="Annotate Masks"),
        ],
    outputs=gr.components.Image(type="numpy", label="Output Image"),
    title="Car Damage Severity Detector",
    examples=SAMPLE["car-dam-sev-det"],
    cache_examples=False,
)

gr.TabbedInterface(
    [car_part_ui, car_dam_ui, car_sev_ui],
    tab_names=[
        'car part segmentation', 
        'car damage detection', 
        'car damage severity detection'
        ]
).queue().launch()