Update README.md

README.md

@@ -98,7 +98,31 @@ Users (both direct and downstream) should be made aware of the risks, biases and
 Use the code below to get started with the model.
 ```python
 from transformers import AutoModelForObjectDetection, AutoImageProcessor
+import torch
+
+image_path=YOUR_IMAGE_PATH
+image = cv2.imread(image_path)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model = AutoModelForObjectDetection.from_pretrained("Yifeng-Liu/rt-detr-finetuned-for-satellite-image-roofs-detection")
 image_processor = AutoImageProcessor.from_pretrained("Yifeng-Liu/rt-detr-finetuned-for-satellite-image-roofs-detection")
+
+
+CONFIDENCE_TRESHOLD = 0.5
+
+with torch.no_grad():
+    model.to(device)
+
+    # load image and predict
+    inputs = image_processor(images=image, return_tensors='pt').to(device)
+    outputs = model(**inputs)
+
+    # post-process
+    target_sizes = torch.tensor([image.shape[:2]]).to(device)
+    results = image_processor.post_process_object_detection(
+        outputs=outputs,
+        threshold=CONFIDENCE_TRESHOLD,
+        target_sizes=target_sizes
+    )[0]
 ```