Spaces:

Sartc
/

tbi

Running

Sartc commited on Mar 25

Commit

3e855cc

verified ·

1 Parent(s): a0d8b60

Update inference.py

Files changed (1) hide show

inference.py CHANGED Viewed

@@ -30,8 +30,9 @@ def save_output(mask, save_path):
     mask_image = Image.fromarray(mask[0])
     mask_image.save(save_path)
-weights_path = "unet_model.pth"
-model = load_model(weights_path, device)
-image_tensor = preprocess_image("DUTS-TE-Image/ILSVRC2012_test_00000003.jpg")
-mask = predict(model, image_tensor, device)
-save_output(mask, "predicted_mask.jpg")

     mask_image = Image.fromarray(mask[0])
     mask_image.save(save_path)
+if __name__ == "__main__":
+    weights_path = "unet_model.pth"
+    model = load_model(weights_path, device)
+    image_tensor = preprocess_image("DUTS-TE-Image/ILSVRC2012_test_00000003.jpg")
+    mask = predict(model, image_tensor, device)
+    save_output(mask, "predicted_mask.jpg")