cafeai commited on
Commit
d3119f5
1 Parent(s): ceb2c24
Files changed (2) hide show
  1. app.py +22 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+
4
+ pipe_aesthetic = pipeline("image-classification", "cafeai/cafe_aesthetic")
5
+ def aesthetic(input_img):
6
+ data = pipe_aesthetic(input_img, top_k=2)
7
+ final = {}
8
+ for d in data:
9
+ final[d["label"]] = d["score"]
10
+ return final
11
+ demo_aesthetic = gr.Interface(fn=aesthetic, inputs=gr.Image(type="pil"), outputs=gr.Label(label="aesthetic"))
12
+
13
+ pipe_style = pipeline("image-classification", "cafeai/cafe_style")
14
+ def style(input_img):
15
+ data = pipe_style(input_img, top_k=5)
16
+ final = {}
17
+ for d in data:
18
+ final[d["label"]] = d["score"]
19
+ return final
20
+ demo_style = gr.Interface(fn=style, inputs=gr.Image(type="pil"), outputs=gr.Label(label="style"))
21
+
22
+ gr.Parallel(demo_aesthetic, demo_style).launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
+ torch
3
+ transformers==4.23.1