Update app.py

app.py

@@ -1,5 +1,4 @@
 import pathlib
-from constants import MODELS_REPO, MODELS_NAMES, MODELS_DETAILS
 
 import gradio as gr
 import torch
@@ -25,12 +24,12 @@ def make_prediction(img, feature_extractor, model):
     )
 
 
-def detect_objects(model_name, image_input, threshold, display_mask=False, img_input_mask=None):
-    feature_extractor = AutoFeatureExtractor.from_pretrained(MODELS_REPO[model_name])
-
-    if "DETR" in model_name:
-        model = DetrForObjectDetection.from_pretrained(MODELS_REPO[model_name])
-        model_details_json = gr.Markdown(MODELS_DETAILS[model_name])
+def detect_objects(model_name, attention_heads_num, image_input, threshold, display_mask=False, img_input_mask=None):
+    model_repo = f"polejowska/detr-r50-cd45rb-all-{str(attention_heads_num)}ah"
+    
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model_repo)
+    
+    model = DetrForObjectDetection.from_pretrained(model_repo)
 
     (
         processed_outputs,
@@ -57,7 +56,6 @@ def detect_objects(model_name, image_input, threshold, display_mask=False, img_i
         viz_img,
         decoder_attention_map_img,
         encoder_attention_map_img,
-        # model_details_json
     )
 
 
@@ -88,14 +86,14 @@ with gr.Blocks(css=css) as app:
                         )
                 with gr.Column():
                     with gr.Row():
-                        options = gr.Dropdown(
+                        model_name = gr.Dropdown(
                             value=MODELS_NAMES[0],
                             choices=MODELS_NAMES,
                             label="Select an object detection model",
                             show_label=True,
                         )
                     with gr.Row():
-                        heads = gr.Dropdown(
+                        attention_heads_num = gr.Dropdown(
                             value=8,
                             choices=[2, 4, 8, 16],
                             label="The number of attention heads in encoder and decoder",
@@ -137,13 +135,11 @@ with gr.Blocks(css=css) as app:
 
     detect_button.click(
         detect_objects,
-        inputs=[options, img_input, slider_input, display_mask, img_input_mask],
+        inputs=[model_name, attention_heads_num, img_input, slider_input, display_mask, img_input_mask],
         outputs=[
             img_output_from_upload,
             decoder_att_map_output,
             encoder_att_map_output,
-            # cross_att_map_output,
-            # model_details,
         ],
         queue=True,
     )