From f7d224f9e4fb8ed30eb75f89063a2baef932aebb Mon Sep 17 00:00:00 2001
From: HaiyangP <46739135+HaiyangPeng@users.noreply.github.com>
Date: Thu, 2 Sep 2021 19:35:30 +0800
Subject: [PATCH] Add the coordinate clipping op before nms (#704)

* add the coordinate clipping op before nms

* modified the pytorch related codes

* modified the pytorch related codes
---
 yolov5/yolov5_trt.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/yolov5/yolov5_trt.py b/yolov5/yolov5_trt.py
index dda55e4..2f655f1 100644
--- a/yolov5/yolov5_trt.py
+++ b/yolov5/yolov5_trt.py
@@ -351,6 +351,11 @@ class YoLov5TRT(object):
         boxes = prediction[prediction[:, 4] >= conf_thres]
         # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]
         boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4])
+        # clip the coordinates
+        boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w -1)
+        boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w -1)
+        boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h -1)
+        boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h -1)
         # Object confidence
         confs = boxes[:, 4]
         # Sort by the confs