Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import PIL.Image as Image
-from ultralytics import YOLO
+from ultralytics import YOLO, RTDETR
 import spaces
 import os
 from huggingface_hub import hf_hub_download
@@ -16,7 +16,11 @@ def get_model_path(model_name):
 @spaces.GPU
 def yolo_inference(images, model_id, conf_threshold, iou_threshold, max_detection):
     model_path = get_model_path(model_id)  # Download model
-    model = YOLO(model_path)
+    if 'rtdetr' in model_id.lower():
+        model_type = RTDETR
+    else:
+        model_type = YOLO
+    model = model_type(model_path)
     results = model.predict(
         source=images,
         conf=conf_threshold,