Update app.py

app.py

@@ -38,6 +38,7 @@ def detect_objects(model_name, image_input, threshold):
 
     if "DETR" in model_name:
         model = DetrForObjectDetection.from_pretrained(MODELS_REPO[model_name])
+        model_details = "DETR details"
     elif "YOLOS" in model_name:
         model = YolosForObjectDetection.from_pretrained(MODELS_REPO[model_name])
 
@@ -64,9 +65,7 @@ def detect_objects(model_name, image_input, threshold):
         decoder_attention_map_img,
         encoder_attention_map_img,
         cross_attention_map_img,
-        gr.outputs.Image(type='file', label=f'attention_head_1'),
-        gr.outputs.Image(type='file', label=f'attention_head_2'),
-        gr.outputs.Image(type='file', label=f'attention_head_3'),
+        model_details
     )
 
 
@@ -121,6 +120,9 @@ with gr.Blocks(css=css) as app:
                 with gr.Column():
                     gr.Markdown("""Cross attentions""")
                     cross_att_map_output = gr.Image(shape=(850, 850))
+        with gr.TabItem("Model details"):
+            with gr.Row():
+                model_details = gr.Markdown(""" """)
 
     detect_button.click(
         detect_objects,
@@ -130,6 +132,7 @@ with gr.Blocks(css=css) as app:
             decoder_att_map_output,
             encoder_att_map_output,
             cross_att_map_output,
+            model_details,
         ],
         queue=True,
     )