Update app.py

app.py

@@ -27,10 +27,11 @@ def make_prediction(img, feature_extractor, model):
 def detect_objects(model_name, attention_heads_num, image_input, threshold, display_mask=False, img_input_mask=None):
     if "CONDITIONAL-DETR" in model_name:
         model_repo = f"polejowska/cdetr-r50-cd45rb-all-4ah"
-        model = DetrForObjectDetection.from_pretrained(model_repo)
+        model = ConditionalDetrForObjectDetection.from_pretrained(model_repo)
     else:
         model_repo = f"polejowska/detr-r50-cd45rb-all-{str(attention_heads_num)}ah"
-        model = ConditionalDetrForObjectDetection.from_pretrained(model_repo)
+        model = DetrForObjectDetection.from_pretrained(model_repo)
+
     
     feature_extractor = AutoFeatureExtractor.from_pretrained(model_repo)