cc1234 commited on
Commit
40613f2
1 Parent(s): bbe49e5

Add skip and translate functionality to predict function

Browse files
Files changed (1) hide show
  1. app.py +15 -0
app.py CHANGED
@@ -46,6 +46,16 @@ def load_labels() -> list[str]:
46
  model = load_model()
47
  labels = load_labels()
48
 
 
 
 
 
 
 
 
 
 
 
49
 
50
  def predict(
51
  image: PIL.Image.Image, score_threshold: float
@@ -68,6 +78,11 @@ def predict(
68
  for index in indices:
69
  label = labels[index]
70
  prob = probs[index]
 
 
 
 
 
71
  result_all[label] = prob
72
  if prob < score_threshold:
73
  break
 
46
  model = load_model()
47
  labels = load_labels()
48
 
49
+ skip = ["rating:safe",
50
+ "rating:questionable",
51
+ "rating:explicit",
52
+ "3d",
53
+ "photorealistic",
54
+ "realistic",
55
+ "uncensored"]
56
+
57
+ translate = {'yuri': 'lesbian'}
58
+
59
 
60
  def predict(
61
  image: PIL.Image.Image, score_threshold: float
 
78
  for index in indices:
79
  label = labels[index]
80
  prob = probs[index]
81
+ if label in skip:
82
+ continue
83
+ if label in translate:
84
+ label = translate[label]
85
+
86
  result_all[label] = prob
87
  if prob < score_threshold:
88
  break