ttheland commited on
Commit
eac1956
1 Parent(s): ac92c82

added no class detected case

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -113,7 +113,11 @@ def do_inference(img):
113
  probs = probs.cpu().numpy()[0]
114
  probs = probs[output]
115
  labels = np.array(class_names)[output]
116
- return {id_to_name[labels[i]]: round(float(probs[i]),2) for i in range(len(labels))}
 
 
 
 
117
 
118
  im = gr.inputs.Image(shape=(512, 512), image_mode='RGB',
119
  invert_colors=False, source="upload",
 
113
  probs = probs.cpu().numpy()[0]
114
  probs = probs[output]
115
  labels = np.array(class_names)[output]
116
+
117
+ if(probs[0] < 0.5):
118
+ return {"No butterfly":1.0}
119
+ else:
120
+ return {id_to_name[labels[i]]: round(float(probs[i]),2) for i in range(len(labels))}
121
 
122
  im = gr.inputs.Image(shape=(512, 512), image_mode='RGB',
123
  invert_colors=False, source="upload",