Show image and heatmap at the same time

app.py

@@ -171,6 +171,7 @@ def infer(url, mp_text, hidx):
     encoded = processor.tokenizer(inferred_token)
     output = processor.decode(encoded.input_ids[0], skip_special_tokens=True)
     selected_token = ''
+    result = None
 
     if hidx > 0 and hidx < len(encoded["input_ids"][0][:-1]):
         input_ids = torch.tensor(encoded.input_ids)
@@ -253,13 +254,13 @@ def infer(url, mp_text, hidx):
         )
         image_rgba = image.copy()
         image_rgba.putalpha(overlay)
-        image = image_rgba
+        result = image_rgba
 
         selected_token = processor.tokenizer.convert_ids_to_tokens(
             encoded["input_ids"][0][hidx]
         )
 
-    return [np.array(image), output, selected_token]
+    return [np.array(image), output, selected_token, result]
 
 
 title = "What's in the picture ?"
@@ -289,6 +290,7 @@ outputs_interface = [
     gr.outputs.Image(label="Image"),
     gr.outputs.Textbox(label="description"),
     gr.outputs.Textbox(label="selected token"),
+    gr.outputs.Image(label="Heatmap")
 ]
 
 interface = gr.Interface(