update

RT-DETR-L_wired_table_cell_det/src/example.py

@@ -1,19 +1,20 @@
 import sys
-sys.dont_write_bytecode = True
-
 import cv2
 import numpy
 
+sys.dont_write_bytecode = True
+
+# Source
 from helper import onnxSessionBuild
 
 pathModel = "./RT-DETR-L_wired_table_cell_det/"
 
-imageSizeDetection = 640
+imageSizeModel = 640
 
 scoreThreshold = 0.3
-
 levelBoxNms = 0.5
 levelBoxContained = 0.9
+countContainedMinimum = 2
 
 onnxSession = onnxSessionBuild(f"{pathModel}onnx/rt-detr-l_wired_table_cell_det.onnx")
 
@@ -27,13 +28,15 @@ def overlapCalculate(coordinateFirstList, coordinateSecondList):
     y2 = min(coordinateFirstList[3], coordinateSecondList[3])
 
     if x2 <= x1 or y2 <= y1:
-        return {"intersectionOverUnion": 0.0, "containment": 0.0}
+        return {"intersectionOverUnion": 0.0, "containment": 0.0, "areaFirst": areaFirst, "areaSecond": areaSecond}
 
     areaIntersection = (x2 - x1) * (y2 - y1)
 
     return {
         "intersectionOverUnion": areaIntersection / float(areaFirst + areaSecond - areaIntersection),
-        "containment": areaIntersection / float(min(areaFirst, areaSecond))
+        "containment": areaIntersection / float(min(areaFirst, areaSecond)),
+        "areaFirst": areaFirst,
+        "areaSecond": areaSecond
     }
 
 def suppressionNonMaximum(itemList):
@@ -47,10 +50,7 @@ def suppressionNonMaximum(itemList):
         for b in range(len(resultList)):
             overlapObject = overlapCalculate(itemSortedList[a]["coordinate"], resultList[b]["coordinate"])
 
-            isSameRegion = overlapObject["intersectionOverUnion"] >= levelBoxNms
-            isInside = overlapObject["containment"] >= levelBoxContained
-
-            if isSameRegion or isInside:
+            if overlapObject["intersectionOverUnion"] >= levelBoxNms:
                 isKeep = False
 
                 break
@@ -60,18 +60,38 @@ def suppressionNonMaximum(itemList):
 
     return resultList
 
+def suppressionContainer(itemList):
+    resultList = []
+
+    for a in range(len(itemList)):
+        countContained = 0
+
+        for b in range(len(itemList)):
+            if a == b:
+                continue
+
+            overlapObject = overlapCalculate(itemList[a]["coordinate"], itemList[b]["coordinate"])
+
+            if overlapObject["containment"] >= levelBoxContained and overlapObject["areaFirst"] > overlapObject["areaSecond"]:
+                countContained += 1
+
+        if countContained < countContainedMinimum:
+            resultList.append(itemList[a])
+
+    return resultList
+
 def inference(imageRgb):
     resultList = []
 
     imageHeight, imageWidth = imageRgb.shape[0:2]
-    imageResized = cv2.resize(imageRgb, (imageSizeDetection, imageSizeDetection), interpolation=cv2.INTER_CUBIC).astype(numpy.float32) / 255.0
+    imageResized = cv2.resize(imageRgb, (imageSizeModel, imageSizeModel), interpolation=cv2.INTER_CUBIC).astype(numpy.float32) / 255.0
 
     tensor = numpy.expand_dims(imageResized.transpose((2, 0, 1)), axis=0).astype(numpy.float32)
 
     tensorFeedObject = {
         "image": tensor,
-        "im_shape": numpy.array([[imageSizeDetection, imageSizeDetection]], dtype=numpy.float32),
-        "scale_factor": numpy.array([[imageSizeDetection / float(imageHeight), imageSizeDetection / float(imageWidth)]], dtype=numpy.float32)
+        "im_shape": numpy.array([[imageSizeModel, imageSizeModel]], dtype=numpy.float32),
+        "scale_factor": numpy.array([[imageSizeModel / float(imageHeight), imageSizeModel / float(imageWidth)]], dtype=numpy.float32)
     }
 
     tensorOutputList = onnxSession.run(None, tensorFeedObject)
@@ -93,7 +113,10 @@ def inference(imageRgb):
                 "coordinate": [x1, y1, x2, y2]
             })
 
-    return suppressionNonMaximum(resultList)
+    resultList = suppressionNonMaximum(resultList)
+    resultList = suppressionContainer(resultList)
+
+    return resultList
 
 image = cv2.imread(sys.argv[1])
 imageRgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)