Fix rgb

2 years ago · 038856cb61
parent a414d1f851
commit 038856cb61
1 changed files with 1 additions and 1 deletions
--- 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)