Upload 8 files

.gitignore

@@ -0,0 +1,5 @@
+.idea/
+.vscode/
+__pycache__/
+
+logs/

README.md

@@ -1,12 +1,3 @@
----
-license: mit
-language:
-- en
-pipeline_tag: image-segmentation
-tags:
-- climate
----
-
 # V-BeachNet
 
 This repository contains the official PyTorch implementation for the paper "A New Framework for Quantifying Alongshore Variability of Swash Motion Using Fully Convolutional Networks." V-BeachNet is built upon V-FloodNet.
@@ -19,46 +10,46 @@ Liang, Y., Li, X., Tsai, B., Chen, Q., & Jafari, N. (2023). V-FloodNet: A video
 
 ## Prerequisites
 
-This code is tested on a newly installed Ubuntu 24.04 with default version of Python and Nvidia GPU.
+Install Conda on your Ubuntu 24.04 with default version of Python and Nvidia GPU.
 
 1. Install Anaconda prerequisite (Can also be accessed from [here](https://docs.anaconda.com/anaconda/install/linux/)):
-    ```sh
-    sudo apt update && \
-    sudo apt install libgl1-mesa-dri libegl1 libglu1-mesa libxrandr2 libxss1 libxcursor1 libxcomposite1 libasound2-data libasound2-plugins libxi6 libxtst6
-    ```
+```sh
+sudo apt update && \
+sudo apt install libgl1-mesa-dri libegl1 libglu1-mesa libxrandr2 libxss1 libxcursor1 libxcomposite1 libasound2-data libasound2-plugins libxi6 libxtst6
+```
 
 2. Download Anaconda3:
-    ```sh
-    curl -O https://repo.anaconda.com/archive/Anaconda3-2024.06-1-Linux-x86_64.sh
-    ```
+```sh
+curl -O https://repo.anaconda.com/archive/Anaconda3-2024.06-1-Linux-x86_64.sh
+```
 
-3. Locate the downloaded file and install it:
-    ```sh
-    bash Anaconda3-2024.06-1-Linux-x86_64.sh
-    ```
+3. Located the downloaded file and install it:
+```sh
+bash Anaconda3-2024.06-1-Linux-x86_64.sh
+```
 
 ## Steps
 
 1. Clone this repository and change directory:
-    ```sh
-    git clone https://huggingface.co/rezasalatin/V-BeachNet.git
-    cd V-BeachNet
-    ```
+```sh
+git clone https://github.com/rezasalatin/V-BeachNet.git
+cd V-BeachNet
+```
 
 2. Create the virtual environment with the requirements:
-    ```sh
-    conda env create -f environment.yml
-    conda activate vbeach
-    ```
-
-3. Visit the "Training_Station" folder and copy your manually segmented (using [labelme](https://github.com/labelmeai/labelme)) dataset to this directory. Open the following file to change any of the variables and save it. Then execute it to train the model:
-    ```sh
-    ./train_video_seg.sh
-    ```
-    Access your trained model from the `log/` directory.
-
-4. Visit the "Testing_Station" folder and copy your data to this directory. Open the following file to change any of the variables (especially the model path from the `log/` folder) and save it. Then execute it to test the model:
-    ```sh
-    ./test_video_seg.sh
-    ```
-    Access your segmented data from the `output` directory.
+```sh
+conda env create -f environment.yml
+conda activate vbeach
+```
+
+3. Visit the "Training_Station" folder and copy your manually segmented dataset to this directory. Open the following file to change any of the variables and save it. Then execute it to train the model:
+```sh
+./train_video_seg.sh
+```
+Access your trained model from log/ directory.
+
+4. Visit the "Testing_Station" folder and copy your data to this directory. Open the following file to change any of the variables (especially model path from the log/ folder) and save it. Then execute it to test the model:
+```sh
+./test_video_seg.sh
+```
+Access your segmented data from output directory.

environment.yml

@@ -0,0 +1,22 @@
+name: vbeach
+channels:
+  - conda-forge
+  - pytorch
+  - nvidia
+dependencies:
+  - tqdm
+  - pytorch
+  - torchvision
+  - torchaudio
+  - pytorch-cuda=12.1
+  - scipy
+  - numpy
+  - matplotlib
+  - pandas
+  - scikit-learn
+  - opencv
+  - gxx_linux-64
+  - pip
+  - pip:
+    - torch-scatter -f https://data.pyg.org/whl/torch-2.1.2+cu121.html
+    - segmentation-models-pytorch==0.2.0

test_image_seg.py

@@ -0,0 +1,173 @@
+import os
+import argparse
+import cv2
+import numpy as np
+
+from PIL import Image
+from pathlib import Path
+from glob import glob
+from tqdm import tqdm
+
+import torch
+import torchvision.transforms as tf
+
+from segmentation_models_pytorch import Linknet
+import myutils
+
+# ROOT_DIR = str(Path(__file__).resolve().parents[0])
+ROOT_DIR = './'
+# time_str = timestr = time.strftime("%Y-%m-%d %H-%M-%S")
+# DEFAULT_OUT = os.path.join(ROOT_DIR, 'output', 'test_waterseg', time_str)
+DEFAULT_OUT = os.path.join(ROOT_DIR, 'output', 'segs')
+# DEFAULT_PALETTE = os.path.join(ROOT_DIR, "assets", "mask_palette.png")
+# sys.path.append(ROOT_DIR)
+# print("Added", ROOT_DIR, "to PATH.")
+
+
+def norm_imagenet(img_pil, dims):
+    """
+    Normalizes and resizes input image
+    :param img_pil: PIL Image
+    :param dims: Model's expected input dimensions
+    :return: Normalized Image as a Tensor
+    """
+
+    # Mean and stddev of ImageNet dataset
+    mean = torch.tensor([0.485, 0.456, 0.406])
+    std = torch.tensor([0.229, 0.224, 0.225])
+
+    # Resize, convert to tensor, normalize
+    transform_norm = tf.Compose([
+        tf.Resize([dims[0], dims[1]]),
+        tf.ToTensor(),
+        tf.Normalize(mean, std)
+    ])
+
+    img_norm = transform_norm(img_pil)
+    return img_norm
+
+
+def predict_one(path, model, mask_outdir, overlay_outdir, device):
+    """
+    Predicts a single image from path
+    :param path: Path to image
+    :param model: Loaded Torch Model
+    :param mask_outdir: Filepath to mask out directory
+    :param overlay_outdir: Filepath to overlay out directory
+    :return: None
+    """
+    img_pil = myutils.load_image_in_PIL(path)
+
+    # Prediction is an PIL Image of 0s and 1s
+    prediction = predict_pil(model, img_pil, model_dims=(416, 416), device=device)
+
+    basename = str(Path(os.path.basename(path)).stem)
+    mask_savepth = os.path.join(mask_outdir, basename + '.png')
+    # mask_save = prediction.convert('RGB')
+    prediction.save(mask_savepth)
+
+    over_savepth = os.path.join(overlay_outdir, basename + '.png')
+    img_np = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
+    overlay_np = myutils.add_overlay(img_np, np.array(prediction))
+    cv2.imwrite(over_savepth, overlay_np)
+    # overlay_np = np.array(img_pil) * 1 + np.array(prediction.convert('RGB')) * 0.8
+    # overlay_np = overlay_np.clip(0, 255)
+    # Image.fromarray(overlay_np.astype(np.uint8)).save(over_savepth)
+
+
+def predict_pil(model, img_pil, model_dims, device):
+    """
+    Predicts a single PIL Image
+    :param model: Loaded PyTorch model
+    :param img_pil: PIL image
+    :param model_dims: Model input dimensions
+    :return: Segmentation prediction as PIL Image
+    """
+
+    img_np = np.array(img_pil)
+    img_tensor_norm = norm_imagenet(img_pil, model_dims)
+
+    # Pipeline to resize the prediction to the original image dimensions
+    pred_resize = tf.Compose([tf.Resize([img_np.shape[0], img_np.shape[1]])])
+
+    # Add extra dimension at front as model expects input 1*3*dimX*dimY (batch size of 1)
+    input_data = img_tensor_norm.unsqueeze(0)
+
+    try:
+        # print("Converted input image to cuda.")
+        prediction = model.predict(input_data.to(device))
+    except:
+        print("Did not convert input image to cuda.")
+        prediction = model.predict(input_data)
+
+    prediction = pred_resize(prediction)
+    prediction = myutils.postprocessing_pred(prediction.squeeze().cpu().round().numpy().astype(np.uint8))
+    prediction = Image.fromarray(prediction).convert('P')
+    prediction.putpalette(myutils.color_palette)
+    return prediction
+
+
+def test_waterseg(model_path, test_path, test_name, out_path, device):
+    """
+    Tests either a single or an entire folder of images
+    :param args: Command line args
+    :return: None
+    """
+    model = torch.load(model_path)
+    print('############################################')
+    print('############################################')
+    print('############################################')
+    test_path = test_path
+    out_path = os.path.join(out_path, test_name)
+
+    mask_out = os.path.join(out_path, 'mask')
+    overlay_out = os.path.join(out_path, 'overlay')
+    if not os.path.exists(mask_out):
+        os.makedirs(mask_out)
+    if not os.path.exists(overlay_out):
+        os.makedirs(overlay_out)
+
+    if os.path.isfile(test_path):
+        predict_one(test_path, model, mask_out, overlay_out, device)
+    elif os.path.isdir(test_path):
+        paths = glob(os.path.join(test_path, '*.jpg')) + glob(os.path.join(test_path, '*.png'))
+        for path in tqdm(paths):
+            predict_one(path, model, mask_out, overlay_out, device)
+    else:
+        print("Error: Unknown path type:", test_path)
+
+
+if __name__ == '__main__':
+    # Hyper parameters
+    parser = argparse.ArgumentParser(description='V-FloodNet: Water Image Segmentation')
+    # Required: Path to the .pth file.
+    parser.add_argument('--model-path',
+                        default='./records/link_efficientb4_model.pth',
+                        type=str,
+                        metavar='PATH',
+                        help='Path to the model')
+    # Required: Path to either the single file or directory of files containing .jpg or .png images
+    parser.add_argument('--test-path',
+                        type=str,
+                        metavar='PATH',
+                        required=True,
+                        help='Can point to folder or an individual jpg/png image')
+    parser.add_argument('--test-name',
+                        type=str,
+                        required=True,
+                        help='Test name')
+    parser.add_argument('--out-path',
+                        default=DEFAULT_OUT,
+                        type=str,
+                        metavar='PATH',
+                        help='(OPTIONAL) Path to output folder, defaults to project root/output')
+    args = parser.parse_args()
+
+    # Device
+    device = torch.device('cpu')
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+
+    test_waterseg(args.model_path, args.test_path, args.test_name, args.out_path, device)
+
+    print(myutils.gct(), 'Test image segmentation done.')

test_video_seg.py

@@ -0,0 +1,132 @@
+import numpy as np
+from tqdm import tqdm, trange
+import os
+import argparse
+from glob import glob
+import torch
+from torch import utils
+from torch.nn import functional as F
+from torchvision.transforms import functional as TF
+from torchvision.transforms import InterpolationMode
+
+from video_module.dataset import Video_DS
+from video_module.model import AFB_URR, FeatureBank
+from test_image_seg import test_waterseg
+import myutils
+
+torch.set_grad_enabled(False)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='V-FloodNet: Water Video Segmentation')
+    parser.add_argument('--gpu', type=int, default=0, help='GPU card id.')
+    parser.add_argument('--budget', type=int, default=250000, help='Max number of features in the feature bank.')
+    parser.add_argument('--viz', action='store_true', default=True, help='Visualize data.')
+    parser.add_argument('--model-path', type=str, required=True, help='Path to the checkpoint.')
+    parser.add_argument('--update-rate', type=float, default=0.1, help='Update Rate for merging new features.')
+    parser.add_argument('--merge-thres', type=float, default=0.95, help='Merging Rate threshold.')
+    parser.add_argument('--test-path', type=str, required=True, help='Path to the test video frames.')
+    parser.add_argument('--test-name', type=str, required=True, help='Name for the test video.')
+    return parser.parse_args()
+
+
+def main(args, device):
+    model = AFB_URR(device, update_bank=True, load_imagenet_params=False)
+    model = model.to(device)
+    model.eval()
+
+    downsample_size = 480
+
+    if os.path.isfile(args.model_path):
+        checkpoint = torch.load(args.model_path)
+        end_epoch = checkpoint['epoch']
+        model.load_state_dict(checkpoint['model'], strict=False)
+        train_loss = checkpoint['loss']
+        seed = checkpoint['seed']
+        print(myutils.gct(),
+              f'Loaded checkpoint {args.model_path}. (end_epoch: {end_epoch}, train_loss: {train_loss}, seed: {seed})')
+    else:
+        print(myutils.gct(), f'No checkpoint found at {args.model_path}')
+        raise IOError
+
+    img_list = sorted(glob(os.path.join(args.test_path, '*.jpg')) + glob(os.path.join(args.test_path, '*.png')))
+    first_frame = myutils.load_image_in_PIL(img_list[0])
+    first_name = os.path.basename(img_list[0])[:-4]
+
+    out_dir = './output/segs'
+    mask_dir = os.path.join(out_dir, args.test_name, 'mask')
+    mask_path = os.path.join(mask_dir, first_name + '.png')
+    if not os.path.exists(mask_path):
+        image_model_path = './records/link_efficientb4_model.pth'
+        test_waterseg(image_model_path, img_list[0], args.test_name, out_dir, device)
+
+    first_mask = myutils.load_image_in_PIL(mask_path, 'P')
+    seq_dataset = Video_DS(img_list, first_frame, first_mask)
+
+    seq_loader = utils.data.DataLoader(seq_dataset, batch_size=1, shuffle=False, num_workers=1)
+
+    seg_dir = os.path.join(out_dir, args.test_name, 'mask')
+    os.makedirs(seg_dir, exist_ok=True)
+    if args.viz:
+        overlay_dir = os.path.join(out_dir, args.test_name, 'overlay')
+        os.makedirs(overlay_dir, exist_ok=True)
+
+    obj_n = seq_dataset.obj_n
+    fb = FeatureBank(obj_n, args.budget, device, update_rate=args.update_rate, thres_close=args.merge_thres)
+
+    ori_first_frame = seq_dataset.first_frame.unsqueeze(0).to(device)
+    ori_first_mask = seq_dataset.first_mask.unsqueeze(0).to(device)
+
+    first_frame = TF.resize(ori_first_frame, downsample_size, InterpolationMode.BICUBIC)
+    first_mask = TF.resize(ori_first_mask, downsample_size, InterpolationMode.NEAREST)
+
+    pred = torch.argmax(ori_first_mask[0], dim=0).cpu().numpy().astype(np.uint8)
+    seg_path = os.path.join(seg_dir, f'{first_name}.png')
+    myutils.save_seg_mask(pred, seg_path, myutils.color_palette)
+
+    if args.viz:
+        overlay_path = os.path.join(overlay_dir, f'{first_name}.png')
+        myutils.save_overlay(ori_first_frame[0], pred, overlay_path, myutils.color_palette)
+
+    with torch.no_grad():
+        k4_list, v4_list = model.memorize(first_frame, first_mask)
+        fb.init_bank(k4_list, v4_list)
+
+        for idx, (frame, frame_name) in enumerate(tqdm(seq_loader)):
+
+            ori_frame = frame.to(device)
+            ori_size = ori_frame.shape[-2:]
+            frame = TF.resize(ori_frame, downsample_size, InterpolationMode.BICUBIC)
+            score, _ = model.segment(frame, fb)
+            pred_mask = F.softmax(score, dim=1)
+
+            k4_list, v4_list = model.memorize(frame, pred_mask)
+            fb.update(k4_list, v4_list, idx + 1)
+
+            pred = TF.resize(pred_mask, ori_size, InterpolationMode.BICUBIC)
+            pred = torch.argmax(pred[0], dim=0).cpu().numpy().astype(np.uint8)
+            pred = myutils.postprocessing_pred(pred)
+            seg_path = os.path.join(seg_dir, f'{frame_name[0]}.png')
+            myutils.save_seg_mask(pred, seg_path, myutils.color_palette)
+            if args.viz:
+                overlay_path = os.path.join(overlay_dir, f'{frame_name[0]}.png')
+                myutils.save_overlay(ori_frame[0], pred, overlay_path, myutils.color_palette)
+
+    fb.print_peak_mem()
+
+
+if __name__ == '__main__':
+
+    args = get_args()
+    print(myutils.gct(), 'Args =', args)
+
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = torch.device('cuda', args.gpu)
+    else:
+        raise ValueError('CUDA is required. --gpu must be >= 0.')
+
+    assert os.path.isdir(args.test_path)
+
+    main(args, device)
+
+    print(myutils.gct(), 'Test video segmentation done.')

test_video_seg.sh

@@ -0,0 +1,4 @@
+python test_video_seg.py \
+	--test-path=Testing_Station/Duck_Rectified/20211017_1400_UTC/ \
+	--test-name=Duck_Rectified \
+	--model-path=logs/20240724-104343/model/best.pth

train_video_seg.py

@@ -0,0 +1,143 @@
+import os
+import time
+import argparse
+import numpy as np
+from tqdm import tqdm
+from glob import glob
+
+import torch
+from torch.utils.data import DataLoader
+
+from video_module.dataset import Water_Image_Train_DS
+from video_module.model import AFB_URR, FeatureBank
+import myutils
+
+# Enable CUDA launch blocking for debugging; set to '0' to disable.
+os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
+
+def get_args():
+    parser = argparse.ArgumentParser(description='Train V-BeachNet')
+    parser.add_argument('--gpu', type=int, default=0, help='GPU card id.')
+    parser.add_argument('--dataset', type=str, required=True, help='Dataset folder.')
+    parser.add_argument('--seed', type=int, default=-1, help='Random seed.')
+    parser.add_argument('--log', action='store_true', help='Save the training results.')
+    parser.add_argument('--lr', type=float, default=1e-5, help='Learning rate (default: 1e-5).')
+    parser.add_argument('--lu', type=float, default=0.5, help='Regularization factor (default: 0.5).')
+    parser.add_argument('--resume', type=str, help='Path to checkpoint (default: none).')
+    parser.add_argument('--new', action='store_true', help='Train the model from scratch.')
+    parser.add_argument('--scheduler-step', type=int, default=25, help='Scheduler step size (default: 25).')
+    parser.add_argument('--total-epochs', type=int, default=100, help='Total number of epochs (default: 100).')
+    parser.add_argument('--budget', type=int, default=300000, help='Maximum number of features in the feature bank (default: 300000).')
+    parser.add_argument('--obj-n', type=int, default=2, help='Maximum number of objects trained simultaneously.')
+    parser.add_argument('--clip-n', type=int, default=6, help='Maximum number of frames in a batch.')
+
+    return parser.parse_args()
+
+def train_model(model, dataloader, criterion, optimizer):
+    stats = myutils.AvgMeter()
+    uncertainty_stats = myutils.AvgMeter()
+    progress_bar = tqdm(dataloader)
+    
+    for _, sample in enumerate(progress_bar):
+        frames, masks, obj_n, info = sample
+
+        if obj_n.item() == 1:
+            continue
+
+        frames, masks = frames[0].to(device), masks[0].to(device)
+        fb_global = FeatureBank(obj_n.item(), args.budget, device)
+        k4_list, v4_list = model.memorize(frames[:1], masks[:1])
+        fb_global.init_bank(k4_list, v4_list)
+
+        scores, uncertainty = model.segment(frames[1:], fb_global)
+        label = torch.argmax(masks[1:], dim=1).long()
+
+        optimizer.zero_grad()
+        loss = criterion(scores, label) + args.lu * uncertainty
+        loss.backward()
+        optimizer.step()
+
+        uncertainty_stats.update(uncertainty.item())
+        stats.update(loss.item())
+        progress_bar.set_postfix(loss=f'{loss.item():.5f} (Avg: {stats.avg:.5f}, Uncertainty Avg: {uncertainty_stats.avg:.5f})')
+
+    return stats.avg
+
+def main():
+    dataset = Water_Image_Train_DS(root=args.dataset, output_size=400, clip_n=args.clip_n, max_obj_n=args.obj_n)
+    dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2, pin_memory=True)
+    print(myutils.gct(), f'Dataset with {len(dataset)} training cases.')
+
+    model = AFB_URR(device, update_bank=False, load_imagenet_params=True).to(device)
+    model.train()
+    model.apply(myutils.set_bn_eval)
+
+    optimizer = torch.optim.AdamW(filter(lambda x: x.requires_grad, model.parameters()), args.lr)
+
+    start_epoch, best_loss = 0, float('inf')
+    if args.resume:
+        if os.path.isfile(args.resume):
+            checkpoint = torch.load(args.resume)
+            model.load_state_dict(checkpoint['model'], strict=False)
+            seed = checkpoint.get('seed', int(time.time()))
+
+            if not args.new:
+                start_epoch = checkpoint['epoch'] + 1
+                optimizer.load_state_dict(checkpoint['optimizer'])
+                best_loss = checkpoint['loss']
+                print(myutils.gct(), f'Resumed from checkpoint {args.resume} (Epoch: {start_epoch-1}, Best Loss: {best_loss}).')
+            else:
+                print(myutils.gct(), f'Loaded checkpoint {args.resume}. Training from scratch.')
+        else:
+            raise FileNotFoundError(f'No checkpoint found at {args.resume}')
+    else:
+        seed = args.seed if args.seed >= 0 else int(time.time())
+
+    print(myutils.gct(), 'Random seed:', seed)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    criterion = torch.nn.CrossEntropyLoss().to(device)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.scheduler_step, gamma=0.5, last_epoch=start_epoch-1)
+
+    for epoch in range(start_epoch, args.total_epochs):
+        print(f'\n{myutils.gct()} Epoch: {epoch}, Learning Rate: {scheduler.get_last_lr()[0]:.6f}')
+        loss = train_model(model, dataloader, criterion, optimizer)
+
+        if args.log:
+            checkpoint = {
+                'epoch': epoch,
+                'model': model.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'loss': loss,
+                'seed': seed
+            }
+            torch.save(checkpoint, os.path.join(model_path, 'final.pth'))
+            if loss < best_loss:
+                best_loss = loss
+                torch.save(checkpoint, os.path.join(model_path, 'best.pth'))
+                print('Updated best model.')
+
+        scheduler.step()
+
+if __name__ == '__main__':
+    args = get_args()
+    print(myutils.gct(), f'Arguments: {args}')
+
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = torch.device('cuda', args.gpu)
+    else:
+        raise ValueError('CUDA is required. Ensure --gpu is set to >= 0.')
+
+    if args.log:
+        log_dir = os.path.join('logs', time.strftime('%Y%m%d-%H%M%S'))
+        model_path = os.path.join(log_dir, 'model')
+        os.makedirs(model_path, exist_ok=True)
+        myutils.save_scripts(log_dir, scripts_to_save=glob('*.*'))
+        myutils.save_scripts(log_dir, scripts_to_save=glob('dataset/*.py', recursive=True))
+        myutils.save_scripts(log_dir, scripts_to_save=glob('model/*.py', recursive=True))
+        myutils.save_scripts(log_dir, scripts_to_save=glob('myutils/*.py', recursive=True))
+        print(myutils.gct(), f'Created log directory: {log_dir}')
+
+    main()
+    print(myutils.gct(), 'Training completed.')

train_video_seg.sh

@@ -0,0 +1,15 @@
+python3 train_video_seg.py \
+    --gpu=0 \
+    --dataset=Training_Station/Duck_Rectified/ \
+    --seed=-1 \
+    --log \
+    --lr=1e-5 \
+    --lu=0.5 \
+    --scheduler-step=25 \
+    --total-epochs=100 \
+    --budget=300000 \
+    --obj-n=2 \
+    --clip-n=6 \
+	--new
+#    --resume=logs/bestmodel/model/best.pth
+# activate new or resume