modify predict.py

modify README.md
yahuiwei123 · Dec 20, 2023 · 14625d9 · 14625d9
1 parent 0b90273
commit 14625d9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -13,7 +13,9 @@ python train.py --batch_size 8 --step 8 --learning_rate 0.01 --num_epochs 100 -o
 ``
 
 ##### 模型预测
-
+``
+python predict.py --image_path './test/img' --step 8 --output_size (128, 128) --output_dir './test/out'
+``
 #### 模型细节
 + 对模型中的如下模块进行了替换
   + Conv2D $\rightarrow$ LayerWiseConvModule

diff --git a/predict.py b/predict.py
@@ -41,6 +41,6 @@
             softmax = nn.Softmax(dim=1)
             logits = softmax(logits)
             mask = torch.argmax(logits, dim=1)
-            cv2.imwrite(output_dir, mask)
+            cv2.imwrite(f"{output_dir}/mask.png", mask)