update gr.outputs.Label

gradio_artist_classifier.py

@@ -98,8 +98,8 @@ def predict(input_image):
         ax3.imshow(t_img)
 
         ax1.set_title(f'Input Image', ha='left', x=0, y=1.05)
-        ax2.set_title(f'Artist Prediction:\n =>{a_label} ({a_prob:.2f})', ha='left', x=0, y=1.05)
-        ax3.set_title(f'Trend Prediction:\n =>{t_label} ({t_prob:.2f})', ha='left', x=0, y=1.05)
+        ax2.set_title(f'Artist Prediction:\n => {a_label} ({a_prob:.2f})', ha='left', x=0, y=1.05)
+        ax3.set_title(f'Style Prediction:\n => {t_label} ({t_prob:.2f})', ha='left', x=0, y=1.05)
         fig.tight_layout()
         
     buf = io.BytesIO()
@@ -107,18 +107,23 @@ def predict(input_image):
     buf.seek(0)
     pil_img = Image.open(buf)
     plt.close()
-    logger.info('--- output generated')
-    return pil_img
+    logger.info('--- image generated')
+
+    a_labels = {id2artist[i]: float(pred) for i, pred in enumerate(a_pred_out)}
+    t_labels = {id2trend[i]: float(pred) for i, pred in enumerate(t_pred_out)}
+    return a_labels, t_labels, pil_img
 
 iface = gr.Interface(
     predict,
     title='Predict Artist and Artistic Style of Drawings 🎨👨🏻‍🎨 (prototype)',
-    description='Upload a drawing and the model will predict how likely it seems given 10 artists and their trend/style',
+    description='Upload a drawing/image and the model will predict how likely it seems given 10 artists and their trend/style',
     inputs=[
         gr.inputs.Image(label='Upload a drawing/image', type='file')
     ],
     outputs=[
-        gr.outputs.Image(label='Prediction')
+        gr.outputs.Label(label='Artists', num_top_classes=5, type='auto'),
+        gr.outputs.Label(label='Styles', num_top_classes=5, type='auto'),
+        gr.outputs.Image(label='Prediction with GradCAM')
     ],
     examples=EXAMPLES,
 )