From 2f8dbe9f6119a5025315108755227479c6f520b2 Mon Sep 17 00:00:00 2001
From: Bhavay Malhotra <rockwizard2001@gmail.com>
Date: Wed, 31 Jul 2024 13:04:06 +0530
Subject: [PATCH 1/2] Added RTDETR files

---
 inference/models/rtdetr/__init__.py |  1 +
 inference/models/rtdetr/rtdetr.py   | 68 +++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 inference/models/rtdetr/__init__.py
 create mode 100644 inference/models/rtdetr/rtdetr.py

diff --git a/inference/models/rtdetr/__init__.py b/inference/models/rtdetr/__init__.py
new file mode 100644
index 000000000..12bcb3fdb
--- /dev/null
+++ b/inference/models/rtdetr/__init__.py
@@ -0,0 +1 @@
+from inference.models.rtdetr.rtdetr import RTDETR
\ No newline at end of file
diff --git a/inference/models/rtdetr/rtdetr.py b/inference/models/rtdetr/rtdetr.py
new file mode 100644
index 000000000..f43b0b17e
--- /dev/null
+++ b/inference/models/rtdetr/rtdetr.py
@@ -0,0 +1,68 @@
+import os
+
+import torch
+import requests
+
+from typing import Any, Tuple
+
+from PIL import Image
+from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
+
+from inference.core.models.base import PreprocessReturnMetadata
+from inference.core.models.roboflow import RoboflowCoreModel
+from inference.core.utils.image_utils import load_image_rgb
+
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+class RTDETR(RoboflowCoreModel):
+    def __init__(self, *args, model_id=f"", **kwargs):
+        super().__init__(*args, model_id=model_id, **kwargs)
+        self.model_id = model_id
+        self.endpoint = model_id
+        self.api_key = API_KEY
+        self.dataset_id, self.version_id = model_id.split("/")
+        self.cache_dir = os.path.join(MODEL_CACHE_DIR, self.endpoint + "/") # "PekingU/rtdetr_r50vd"
+        dtype = torch.bfloat16
+        self.model = RTDetrForObjectDetection.from_pretrained(
+            self.cache_dir,
+            torch_dtype=dtype,
+            device_map=DEVICE,
+            revision="bfloat16",
+        ).eval()
+
+        self.processor = RTDetrImageProcessor.from_pretrained(
+            self.cache_dir,
+        )
+        # self.task_type = "lmm"
+
+    def preprocess(
+        self, image: Any, **kwargs
+    ) -> Tuple[Image.Image, PreprocessReturnMetadata]:
+        pil_image = Image.fromarray(load_image_rgb(image))
+
+        return pil_image, PreprocessReturnMetadata({})
+
+    def postprocess(
+        self,
+        predictions: Tuple[str],
+        preprocess_return_metadata: PreprocessReturnMetadata,
+        **kwargs,
+    ) -> Any:
+        return predictions[0]
+    
+    def predict(self, image_in: Image.Image, **kwargs):
+        model_inputs = self.processor(
+            images=image_in, return_tensors="pt"
+        ).to(self.model.device)
+
+        with torch.inference_mode():
+            outputs = self.model(**model_inputs)
+            
+        results = self.image_processor.post_process_object_detection(outputs, target_sizes=torch.tensor([image_in.size[::-1]]), threshold=0.3)
+
+        return results
+    
+if __name__ == "__main__":
+    m = RTDETR()
+    print(m.infer())
+

From 8545b09472a42df9298fbd986c5a10aa7046cbea Mon Sep 17 00:00:00 2001
From: Bhavay Malhotra <rockwizard2001@gmail.com>
Date: Fri, 2 Aug 2024 20:42:19 +0530
Subject: [PATCH 2/2] Changed base class to TransformerModel

---
 inference/models/rtdetr/rtdetr.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/inference/models/rtdetr/rtdetr.py b/inference/models/rtdetr/rtdetr.py
index f43b0b17e..993e9d10b 100644
--- a/inference/models/rtdetr/rtdetr.py
+++ b/inference/models/rtdetr/rtdetr.py
@@ -9,12 +9,12 @@
 from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
 
 from inference.core.models.base import PreprocessReturnMetadata
-from inference.core.models.roboflow import RoboflowCoreModel
+from inference.models.transformers.transformers import TransformerModel
 from inference.core.utils.image_utils import load_image_rgb
 
 DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-class RTDETR(RoboflowCoreModel):
+class RTDETR(TransformerModel):
     def __init__(self, *args, model_id=f"", **kwargs):
         super().__init__(*args, model_id=model_id, **kwargs)
         self.model_id = model_id