Update app.py

app.py

@@ -1,31 +1,33 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, ImageClassificationPipeline
 
-pipe_aesthetic = pipeline("image-classification", "./sonic")
-def aesthetic(input_img):
-    data = pipe_aesthetic(input_img, top_k=2)
-    final = {}
-    for d in data:
-        final[d["label"]] = d["score"]
-    return final
-demo_aesthetic = gr.Interface(fn=aesthetic, inputs=gr.Image(type="pil"), outputs=gr.Label(label="aesthetic"))
+class MultiClassLabel(ImageClassificationPipeline):
+    def postprocess(self, model_outputs, top_k=5):
+        if top_k > self.model.config.num_labels:
+            top_k = self.model.config.num_labels
 
-pipe_style = pipeline("image-classification", "cafeai/cafe_style")
-def style(input_img):
-    data = pipe_style(input_img, top_k=5)
-    final = {}
-    for d in data:
-        final[d["label"]] = d["score"]
-    return final
-demo_style = gr.Interface(fn=style, inputs=gr.Image(type="pil"), outputs=gr.Label(label="style"))
+        if self.framework == "pt":
+            probs = model_outputs.logits.sigmoid()[0]
+            scores, ids = probs.topk(top_k)
+        elif self.framework == "tf":
+            probs = stable_softmax(model_outputs.logits, axis=-1)[0]
+            topk = tf.math.top_k(probs, k=top_k)
+            scores, ids = topk.values.numpy(), topk.indices.numpy()
+        else:
+            raise ValueError(f"Unsupported framework: {self.framework}")
 
-pipe_waifu = pipeline("image-classification", "cafeai/cafe_waifu")
-def waifu(input_img):
-    data = pipe_waifu(input_img, top_k=5)
+        scores = scores.tolist()
+        ids = ids.tolist()
+        return [{"score": score, "label": self.model.config.id2label[_id]} for score, _id in zip(scores, ids)]
+
+pipe_aesthetic = pipeline("image-classification", "./sonic", pipeline_class=MultiClassLabel)
+
+def aesthetic(input_img):
+    data = pipe_aesthetic(input_img, top_k=5)
     final = {}
     for d in data:
         final[d["label"]] = d["score"]
     return final
-demo_waifu = gr.Interface(fn=waifu, inputs=gr.Image(type="pil"), outputs=gr.Label(label="waifu"))
+demo_aesthetic = gr.Interface(fn=aesthetic, inputs=gr.Image(type="pil"), outputs=gr.Label(label="characters"))
 
-gr.Parallel(demo_aesthetic, demo_style, demo_waifu).launch()
+gr.Parallel(demo_aesthetic).launch()