From 038856cb61b1cc1d7474803e58b60185c434b521 Mon Sep 17 00:00:00 2001
From: triple-Mu <gpu@163.com>
Date: Fri, 13 Jan 2023 13:53:24 +0800
Subject: [PATCH] Fix rgb

---
 infer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/infer.py b/infer.py
index 3b0aa2c..a074af2 100644
--- a/infer.py
+++ b/infer.py
@@ -134,7 +134,7 @@ def main(args):
         data = Engine(tensor)
 
         if args.seg:
-            seg_img = torch.asarray(seg_img[dh:H - dh, dw:W - dw, :],
+            seg_img = torch.asarray(seg_img[dh:H - dh, dw:W - dw, [2, 1, 0]],
                                     device=device)
             bboxes, scores, labels, masks = seg_postprocess(
                 data, bgr.shape[:2], args.conf_thres, args.iou_thres)