diff --git a/.gitignore b/.gitignore
index c7bb7ca..9d60c97 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,8 +1,6 @@
 .venv/
 .mypy_cache/
 __pycache__/
-
-checkpoints/
 wandb/
 
-INTERRUPTED.pth
+*.pth
diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 0000000..f7a9165
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,22 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Python: Current File",
+      "type": "python",
+      "request": "launch",
+      "program": "${file}",
+      "console": "integratedTerminal",
+      "args": [
+        "--input",
+        "SM.png",
+        "--output",
+        "test.png",
+      ],
+      "justMyCode": true
+    }
+  ]
+}
diff --git a/SM.png.REMOVED.git-id b/SM.png.REMOVED.git-id
new file mode 100644
index 0000000..21a947e
--- /dev/null
+++ b/SM.png.REMOVED.git-id
@@ -0,0 +1 @@
+c6d08aa612451072cfe32a3ee086d08342ed9dd9
\ No newline at end of file
diff --git a/src/predict.py b/src/predict.py
index c1a445c..587e11e 100755
--- a/src/predict.py
+++ b/src/predict.py
@@ -1,16 +1,13 @@
 import argparse
 import logging
-import os
 
+import albumentations as A
 import numpy as np
 import torch
-import torch.nn.functional as F
+from albumentations.pytorch import ToTensorV2
 from PIL import Image
-from torchvision import transforms
 
-from src.utils.dataset import BasicDataset
 from unet import UNet
-from utils.utils import plot_img_and_mask
 
 
 def get_args():
@@ -20,7 +17,7 @@ def get_args():
     parser.add_argument(
         "--model",
         "-m",
-        default="MODEL.pth",
+        default="model.pth",
         metavar="FILE",
         help="Specify the file in which the model is stored",
     )
@@ -28,7 +25,6 @@ def get_args():
         "--input",
         "-i",
         metavar="INPUT",
-        nargs="+",
         help="Filenames of input images",
         required=True,
     )
@@ -36,108 +32,62 @@ def get_args():
         "--output",
         "-o",
         metavar="OUTPUT",
-        nargs="+",
         help="Filenames of output images",
     )
     parser.add_argument(
-        "--viz",
-        "-v",
-        action="store_true",
-        help="Visualize the images as they are processed",
-    )
-    parser.add_argument(
-        "--no-save",
-        "-n",
-        action="store_true",
-        help="Do not save the output masks",
-    )
-    parser.add_argument(
-        "--mask-threshold",
+        "--threshold",
         "-t",
         type=float,
         default=0.5,
         help="Minimum probability value to consider a mask pixel white",
     )
-    parser.add_argument(
-        "--scale",
-        "-s",
-        type=float,
-        default=0.5,
-        help="Scale factor for the input images",
-    )
 
     return parser.parse_args()
 
 
-def predict_img(net, full_img, device, scale_factor=1, out_threshold=0.5):
-    net.eval()
-    img = torch.from_numpy(BasicDataset.preprocess(full_img, scale_factor, is_mask=False))
+def predict_img(net, img, device, threshold):
     img = img.unsqueeze(0)
     img = img.to(device=device, dtype=torch.float32)
 
+    net.eval()
     with torch.inference_mode():
         output = net(img)
+        preds = torch.sigmoid(output)[0]
+        full_mask = preds.cpu().squeeze()
 
-        probs = torch.sigmoid(output)[0]
-
-        tf = transforms.Compose(
-            [transforms.ToPILImage(), transforms.Resize((full_img.size[1], full_img.size[0])), transforms.ToTensor()]
-        )
-
-        full_mask = tf(probs.cpu()).squeeze()
-
-    if net.n_classes == 1:
-        return (full_mask > out_threshold).numpy()
-    else:
-        return F.one_hot(full_mask.argmax(dim=0), net.n_classes).permute(2, 0, 1).numpy()
-
-
-def get_output_filenames(args):
-    def _generate_name(fn):
-        split = os.path.splitext(fn)
-        return f"{split[0]}_OUT{split[1]}"
-
-    return args.output or list(map(_generate_name, args.input))
-
-
-def mask_to_image(mask: np.ndarray):
-    if mask.ndim == 2:
-        return Image.fromarray((mask * 255).astype(np.uint8))
-    elif mask.ndim == 3:
-        return Image.fromarray((np.argmax(mask, axis=0) * 255 / mask.shape[0]).astype(np.uint8))
+    return np.asarray(full_mask > threshold)
 
 
 if __name__ == "__main__":
     args = get_args()
-    in_files = args.input
-    out_files = get_output_filenames(args)
 
-    net = UNet(n_channels=3, n_classes=2)
+    net = UNet(n_channels=3, n_classes=1)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    logging.info(f"Loading model {args.model}")
     logging.info(f"Using device {device}")
 
+    logging.info("Transfering model to device")
     net.to(device=device)
+
+    logging.info(f"Loading model {args.model}")
     net.load_state_dict(torch.load(args.model, map_location=device))
 
-    logging.info("Model loaded!")
+    logging.info(f"Loading image {args.input}")
+    img = Image.open(args.input).convert("RGB")
 
-    for i, filename in enumerate(in_files):
+    logging.info(f"Preprocessing image {args.input}")
+    tf = A.Compose(
+        [
+            A.ToFloat(max_value=255),
+            ToTensorV2(),
+        ],
+    )
+    aug = tf(image=np.asarray(img))
+    img = aug["image"]
 
-        logging.info(f"\nPredicting image {filename} ...")
-        img = Image.open(filename)
+    logging.info(f"Predicting image {args.input}")
+    mask = predict_img(net=net, img=img, threshold=args.threshold, device=device)
 
-        mask = predict_img(
-            net=net, full_img=img, scale_factor=args.scale, out_threshold=args.mask_threshold, device=device
-        )
-
-        if not args.no_save:
-            out_filename = out_files[i]
-            result = mask_to_image(mask)
-            result.save(out_filename)
-            logging.info(f"Mask saved to {out_filename}")
-
-        if args.viz:
-            logging.info(f"Visualizing results for image {filename}, close to continue...")
-            plot_img_and_mask(img, mask)
+    logging.info(f"Saving prediction to {args.output}")
+    mask = Image.fromarray(mask)
+    mask.write(args.output)
diff --git a/src/train.py b/src/train.py
index 1fd7468..4f56bde 100644
--- a/src/train.py
+++ b/src/train.py
@@ -105,6 +105,9 @@ def main():
         net.load_state_dict(torch.load(args.load, map_location=device))
         logging.info(f"Model loaded from {args.load}")
 
+    # save initial model.pth
+    torch.save(net.state_dict(), "model.pth")
+
     # transfer network to device
     net.to(device=device)
 
@@ -146,7 +149,7 @@ def main():
     criterion = nn.BCEWithLogitsLoss()
 
     # setup wandb
-    wandb.init(
+    run = wandb.init(
         project="U-Net-tmp",
         config=dict(
             epochs=args.epochs,
@@ -156,6 +159,9 @@ def main():
         ),
     )
     wandb.watch(net, log_freq=100)
+    # artifact = wandb.Artifact("model", type="model")
+    # artifact.add_file("model.pth")
+    # run.log_artifact(artifact)
 
     logging.info(
         f"""Starting training:
@@ -222,9 +228,11 @@ def main():
             print(f"Train Loss: {train_loss:.3f}, Valid Score: {val_score:3f}")
 
             # save weights when epoch end
-            Path(CHECKPOINT_DIR).mkdir(parents=True, exist_ok=True)
-            torch.save(net.state_dict(), str(CHECKPOINT_DIR / "checkpoint_epoch{}.pth".format(epoch)))
-            logging.info(f"Checkpoint {epoch} saved!")
+            # torch.save(net.state_dict(), "model.pth")
+            # run.log_artifact(artifact)
+            logging.info(f"model saved!")
+
+        run.finish()
 
     except KeyboardInterrupt:
         torch.save(net.state_dict(), "INTERRUPTED.pth")
diff --git a/test.png b/test.png
new file mode 100644
index 0000000..a0cad88
Binary files /dev/null and b/test.png differ