cdnuts commited on
Commit
81e3554
1 Parent(s): 9afdd42

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -32,14 +32,14 @@ allowed_tags.append("explicit")
32
  allowed_tags.append("questionable")
33
  allowed_tags.append("safe")
34
 
35
- def create_tags(image):
36
  img = image.convert('RGB')
37
  tensor = transform(img).unsqueeze(0)
38
 
39
  with torch.no_grad():
40
  out = model(tensor)
41
  probabilities = torch.nn.functional.sigmoid(out[0])
42
- indices = torch.where(probabilities > 0.3)[0]
43
  values = probabilities[indices]
44
 
45
  temp = []
@@ -56,7 +56,7 @@ def create_tags(image):
56
 
57
  demo = gr.Interface(
58
  fn=create_tags,
59
- inputs=[gr.Image(type="pil")],
60
  outputs=["text"],
61
  )
62
  demo.launch()
 
32
  allowed_tags.append("questionable")
33
  allowed_tags.append("safe")
34
 
35
+ def create_tags(image, thres):
36
  img = image.convert('RGB')
37
  tensor = transform(img).unsqueeze(0)
38
 
39
  with torch.no_grad():
40
  out = model(tensor)
41
  probabilities = torch.nn.functional.sigmoid(out[0])
42
+ indices = torch.where(probabilities > thres)[0]
43
  values = probabilities[indices]
44
 
45
  temp = []
 
56
 
57
  demo = gr.Interface(
58
  fn=create_tags,
59
+ inputs=[gr.Image(type="pil"), gr.Slider(minimum=0, maximum = 1, step = 0.01, value = 0.3)],
60
  outputs=["text"],
61
  )
62
  demo.launch()