#Bug Fixed: Predict mask does not match with original image size

HengWeiBin · Mar 20, 2023 · 5b51b74 · 5b51b74
1 parent ff9d4d0
commit 5b51b74
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tools/custom.py b/tools/custom.py
@@ -3,7 +3,7 @@
 # Modified based on https://github.com/XuJiacong/PIDNet
 # ------------------------------------------------------------------------------
 # ------------------------------------------------------------------------------
-# example command: python3 tools\custom.py --model-type 'pidnet-l' --model output\elanroad\best.pt --input samples\Cam5.mp4 --n-class 2 --visualize --show
+# example command: python tools\custom.py --model-type 'pidnet-l' --model output\elanroad\pidnet_large_elan_HSV_MBA\best.pt --input samples\road.jpg --n-class 2 --visualize --show
 # ------------------------------------------------------------------------------
 
 import glob
@@ -93,7 +93,7 @@ def main(args):
             t1 = time_synchronized()
             pred = model(img)
             t2 = time_synchronized()
-            pred = F.interpolate(pred, size=img.size()[-2:], 
+            pred = F.interpolate(pred, size=imgOrigin.shape[:-1], 
                                  mode='bilinear', align_corners=True)
             pred = torch.argmax(pred, dim=1).squeeze(0).cpu().numpy()