Update README.md

README.md

@@ -49,6 +49,7 @@ Now, let's compare the performance of this model with other models.
 ### Direct Use
 
 ```python
+import cv2
 import torch
 from transformers import AutoImageProcessor, AutoModel
 
@@ -67,6 +68,28 @@ segmentation = img_proc.post_process_semantic_segmentation(
     output,
     target_sizes=[img.size[::-1]]
 )
+
+def detect_bboxes(masks: np.ndarray):
+    detected_blocks = []
+    contours, _ = cv2.findContours(masks.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    for contour in list(contours):
+        if len(list(contour)) >= 4:
+            # smallest rectangle containing all points in the contour is computed
+            x, y, width, height = cv2.boundingRect(contour)
+            bounding_box = [x, y, x + width, y + height]
+            detected_blocks.append(bounding_box)
+    return detected_blocks
+
+pred_bbox = []
+for segment in segmentation:
+    bboxes, labels = [], []
+    for ii in range(1, len(items)):
+        mm = segment == ii
+        if mm.sum() > 0:
+            bbx, lab = detect_bboxes(mm.numpy())
+            bboxes.extend(bbx)
+            labels.extend([ii]*len(bbx))
+    pred_bbox.append(dict(bboxes=bboxes, labels=lables))
 ```
 
 ### Citation