TheOneReborn's picture
feat: modify interface and add samples
b734e06
raw
history blame
651 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_v1.pkl')
def categorize_image(image):
prediction, index, probabilities = LEARNER.predict(image)
return dict(zip(CATEGORIES, map(float, probabilities)))
print([image for image in TEST_IMAGES_PATH.ls()])
demo = gradio.Interface(
categorize_image,
inputs='image',
outputs='label',
examples=[image for image in TEST_IMAGES_PATH]
)
demo.launch()