TheOneReborn's picture
feat: second model
3e156b5
raw
history blame contribute delete
616 Bytes
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()