import gradio

from pathlib import Path

from fastai.vision.all import (
    load_learner
)
import gradio.interface


CATEGORIES = ('Damaged', 'Whole')
MODEL_PATH = Path('.') / 'models'
TEST_IMAGES_PATH = Path('.') / 'test'
LEARNER = load_learner(MODEL_PATH / 'car-damage-detection_v2.pkl')


def categorize_image(image):
    prediction, index, probabilities = LEARNER.predict(image)
    return dict(zip(CATEGORIES, map(float, probabilities)))


demo = gradio.Interface(
    categorize_image,
    inputs='image',
    outputs='label',
    examples=[str(image) for image in TEST_IMAGES_PATH.iterdir()]
)
demo.launch()