update

RT-DETR-L_wireless_table_cell_det/src/example.py

@@ -0,0 +1,109 @@
+import sys
+sys.dont_write_bytecode = True
+
+import cv2
+import numpy
+
+from helper import onnxSessionBuild
+
+pathModel = "./RT-DETR-L_wireless_table_cell_det/"
+
+imageSizeDetection = 640
+
+scoreThreshold = 0.3
+
+levelBoxNms = 0.5
+levelBoxContained = 0.9
+
+onnxSession = onnxSessionBuild(f"{pathModel}onnx/rt-detr-l_wireless_table_cell_det.onnx")
+
+def overlapCalculate(coordinateFirstList, coordinateSecondList):
+    areaFirst = (coordinateFirstList[2] - coordinateFirstList[0]) * (coordinateFirstList[3] - coordinateFirstList[1])
+    areaSecond = (coordinateSecondList[2] - coordinateSecondList[0]) * (coordinateSecondList[3] - coordinateSecondList[1])
+
+    x1 = max(coordinateFirstList[0], coordinateSecondList[0])
+    y1 = max(coordinateFirstList[1], coordinateSecondList[1])
+    x2 = min(coordinateFirstList[2], coordinateSecondList[2])
+    y2 = min(coordinateFirstList[3], coordinateSecondList[3])
+
+    if x2 <= x1 or y2 <= y1:
+        return {"intersectionOverUnion": 0.0, "containment": 0.0}
+
+    areaIntersection = (x2 - x1) * (y2 - y1)
+
+    return {
+        "intersectionOverUnion": areaIntersection / float(areaFirst + areaSecond - areaIntersection),
+        "containment": areaIntersection / float(min(areaFirst, areaSecond))
+    }
+
+def suppressionNonMaximum(itemList):
+    resultList = []
+
+    itemSortedList = sorted(itemList, key=lambda itemObject: itemObject["score"], reverse=True)
+
+    for a in range(len(itemSortedList)):
+        isKeep = True
+
+        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:
+                isKeep = False
+
+                break
+
+        if isKeep:
+            resultList.append(itemSortedList[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
+
+    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)
+    }
+
+    tensorOutputList = onnxSession.run(None, tensorFeedObject)
+
+    boxCount = int(tensorOutputList[1][0]) if len(tensorOutputList) > 1 else len(tensorOutputList[0])
+
+    for a in range(boxCount):
+        value = tensorOutputList[0][a]
+
+        score = float(value[1])
+        x1 = max(0.0, min(float(value[2]), float(imageWidth)))
+        y1 = max(0.0, min(float(value[3]), float(imageHeight)))
+        x2 = max(0.0, min(float(value[4]), float(imageWidth)))
+        y2 = max(0.0, min(float(value[5]), float(imageHeight)))
+
+        if score >= scoreThreshold and x2 > x1 and y2 > y1:
+            resultList.append({
+                "score": score,
+                "coordinate": [x1, y1, x2, y2]
+            })
+
+    return suppressionNonMaximum(resultList)
+
+image = cv2.imread(sys.argv[1])
+imageRgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+itemList = inference(imageRgb)
+
+for a in range(len(itemList)):
+    coordinateList = []
+
+    for b in range(len(itemList[a]["coordinate"])):
+        coordinateList.append(int(round(itemList[a]["coordinate"][b])))
+
+    print(f"{itemList[a]['score']:.6f} | {coordinateList}")