Only forward pass from trained model file.

madhavkhoslaa · Jul 15, 2019 · 1dae603 · 1dae603
1 parent afe34b5
commit 1dae603
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/predict.py b/predict.py
@@ -0,0 +1,14 @@
+import torch
+import torchvision
+from model.Unet import UNeT
+import sys
+import torchvision
+
+to_pred= sys.argv[1:]
+model= UNeT(n_channels= 3, n_classes=30)
+model.load_state_dict('./model.pt')
+model.eval()
+for _ in to_pred:
+    output= model(_)
+    torchvision.utils.save_image(output, _+'prediction.jpeg')
+    print("Output saved at {}+prediction.jpeg".format(_))