Vahe commited on
Commit
23cfba6
·
1 Parent(s): 7965073

category list derived

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -43,7 +43,8 @@ def main():
43
  model = load_learner(MODEL_PATH)
44
  model_output = model.predict('fruit_image.jpg')
45
 
46
- prob_idx = model.dls.vocab.index(model_output[0])
 
47
  st.write(f'{model_output[0].title()} is depicted in the photo with {model_output[-1][prob_idx]:.4f} confidence.')
48
 
49
  st.session_state.pop("fruit")
 
43
  model = load_learner(MODEL_PATH)
44
  model_output = model.predict('fruit_image.jpg')
45
 
46
+ category_list = [cat for cat in model.dls.vocab]
47
+ prob_idx = category_list.index(model_output[0])
48
  st.write(f'{model_output[0].title()} is depicted in the photo with {model_output[-1][prob_idx]:.4f} confidence.')
49
 
50
  st.session_state.pop("fruit")