radames commited on
Commit
c9267e5
1 Parent(s): 4fecf45
Files changed (1) hide show
  1. app.py +4 -6
app.py CHANGED
@@ -25,7 +25,6 @@ aesthetic_pipe = pipeline("image-classification",
25
  torch_dtype=dtype)
26
 
27
  def predict(image, files=None):
28
- print(image, files)
29
  images_paths = [image]
30
  if not files == None:
31
  images_paths = list(map(lambda x: x.name, files))
@@ -35,12 +34,11 @@ def predict(image, files=None):
35
  aesthetic = aesthetic_pipe(pil_images)
36
  nsfw = nsfw_pipe(pil_images)
37
  results = [ a + b + c for (a,b,c) in zip(style, aesthetic, nsfw)]
38
-
39
  label_data = {}
40
- if image:
41
- label_data = [{ row["label"]:row["score"] for row in image } for image in results]
42
-
43
- return label_data[0], results
44
 
45
  with gr.Blocks() as blocks:
46
  with gr.Row():
 
25
  torch_dtype=dtype)
26
 
27
  def predict(image, files=None):
 
28
  images_paths = [image]
29
  if not files == None:
30
  images_paths = list(map(lambda x: x.name, files))
 
34
  aesthetic = aesthetic_pipe(pil_images)
35
  nsfw = nsfw_pipe(pil_images)
36
  results = [ a + b + c for (a,b,c) in zip(style, aesthetic, nsfw)]
 
37
  label_data = {}
38
+ if image is not None:
39
+ label_data = { row["label"]:row["score"] for row in results[0] }
40
+
41
+ return label_data, results
42
 
43
  with gr.Blocks() as blocks:
44
  with gr.Row():