Update inference.py

inference.py

@@ -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")
+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")