spuun commited on
Commit
7c1e4c2
1 Parent(s): a533c0a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -48,7 +48,7 @@ def launch(img):
48
  tm_image = load_image(img, mode='RGB')
49
  tm_input_ = _img_encode(tm_image, size=(256, 256))[None, ...]
50
  tm_items, = nsfw_tm.run(['output'], {'input': tm_input_})
51
- tm_output = list(zip(tm_cfg["labels"], map(lambda x: x.item(), tm_items[0])))[0][0]
52
 
53
  match tm_output:
54
  case "safe":
 
48
  tm_image = load_image(img, mode='RGB')
49
  tm_input_ = _img_encode(tm_image, size=(256, 256))[None, ...]
50
  tm_items, = nsfw_tm.run(['output'], {'input': tm_input_})
51
+ tm_output = sorted(list(zip(tm_cfg["labels"], map(lambda x: x.item(), tm_items[0]))), key=lambda x: x[1])[0][0]
52
 
53
  match tm_output:
54
  case "safe":