Upload 13 files

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+RCFPyTorch0/vgg16convs.mat filter=lfs diff=lfs merge=lfs -text

RCFPyTorch0

@@ -1 +0,0 @@
-Subproject commit 0f1f2486e5cca2f0c564fc87bdd87b182bfb03c1

RCFPyTorch0/LICENSE.md

@@ -0,0 +1 @@
+The code is released under the [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License](https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode) for NonCommercial use only. Any commercial use should get formal permission first (Email: yun.liu@vision.ee.ethz.ch).

RCFPyTorch0/README.md

@@ -0,0 +1,68 @@
+## [Richer Convolutional Features for Edge Detection](http://mmcheng.net/rcfedge/)
+
+This is the PyTorch implementation of our edge detection method, RCF.
+
+### Citations
+
+If you are using the code/model/data provided here in a publication, please consider citing:
+
+    @article{liu2019richer,
+      title={Richer Convolutional Features for Edge Detection},
+      author={Liu, Yun and Cheng, Ming-Ming and Hu, Xiaowei and Bian, Jia-Wang and Zhang, Le and Bai, Xiang and Tang, Jinhui},
+      journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
+      volume={41},
+      number={8},
+      pages={1939--1946},
+      year={2019},
+      publisher={IEEE}
+    }
+    
+### Training
+
+1. Clone the RCF repository:
+    ```
+    git clone https://github.com/yun-liu/RCF-PyTorch.git
+    ```
+
+2. Download the ImageNet-pretrained model ([Google Drive](https://drive.google.com/file/d/1szqDNG3dUO6BM3l6YBuC9vWp16n48-cK/view?usp=sharing) or [Baidu Yun](https://pan.baidu.com/s/1vfntX-cTKnk58atNW5T1lA?pwd=g5af)), and put it into the `$ROOT_DIR` folder.
+
+3. Download the datasets as below, and extract these datasets to the `$ROOT_DIR/data/` folder.
+
+    ```
+    wget http://mftp.mmcheng.net/liuyun/rcf/data/bsds_pascal_train_pair.lst
+    wget http://mftp.mmcheng.net/liuyun/rcf/data/HED-BSDS.tar.gz
+    wget http://mftp.mmcheng.net/liuyun/rcf/data/PASCAL.tar.gz
+    ```
+    
+4. Run the following command to start the training:
+    ```
+    python train.py --save-dir /path/to/output/directory/
+    ```
+    
+### Testing
+
+1. Download the pretrained model (BSDS500+PASCAL: [Google Drive](https://drive.google.com/file/d/1oxlHQCM4mm5zhHzmE7yho_oToU5Ucckk/view?usp=sharing) or [Baidu Yun](https://pan.baidu.com/s/1Tpf_-dIxHmKwH5IeClt0Ng?pwd=03ad)), and put it into the `$ROOT_DIR` folder.
+
+2. Run the following command to start the testing:
+    ```
+    python test.py --checkpoint bsds500_pascal_model.pth --save-dir /path/to/output/directory/
+    ```
+   This pretrained model should achieve an ODS F-measure of 0.812.
+
+For more information about RCF and edge quality evaluation, please refer to this page: [yun-liu/RCF](https://github.com/yun-liu/RCF)
+
+### Edge PR Curves
+
+We have released the code and data for plotting the edge PR curves of many existing edge detectors [here](https://github.com/yun-liu/plot-edge-pr-curves).
+
+### RCF based on other frameworks 
+
+Caffe based RCF: [yun-liu/RCF](https://github.com/yun-liu/RCF)
+
+Jittor based RCF: [yun-liu/RCF-Jittor](https://github.com/yun-liu/RCF-Jittor)
+
+### Acknowledgements
+
+[1] [balajiselvaraj1601/RCF_Pytorch_Updated](https://github.com/balajiselvaraj1601/RCF_Pytorch_Updated)
+
+[2] [meteorshowers/RCF-pytorch](https://github.com/meteorshowers/RCF-pytorch)

RCFPyTorch0/bsds500_pascal_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9913d9ae1eaa4a71022e89e8c8f6e3eeab5f9bd1cb6a2cc91b1bba7bf36e898c
+size 59235375

RCFPyTorch0/dataset.py

@@ -0,0 +1,45 @@
+import torch
+import cv2
+import numpy as np
+import os.path as osp
+
+
+class BSDS_Dataset(torch.utils.data.Dataset):
+    def __init__(self, root='data/HED-BSDS', split='test', transform=False):
+        super(BSDS_Dataset, self).__init__()
+        self.root = root
+        self.split = split
+        self.transform = transform
+        if self.split == 'train':
+            self.file_list = osp.join(self.root, 'bsds_pascal_train_pair.lst')
+        elif self.split == 'test':
+            self.file_list = osp.join(self.root, 'test.lst')
+        else:
+            raise ValueError('Invalid split type!')
+        with open(self.file_list, 'r') as f:
+            self.file_list = f.readlines()
+        self.mean = np.array([104.00698793, 116.66876762, 122.67891434], dtype=np.float32)
+
+    def __len__(self):
+        return len(self.file_list)
+
+    def __getitem__(self, index):
+        if self.split == 'train':
+            img_file, label_file = self.file_list[index].split()
+            label = cv2.imread(osp.join(self.root, label_file), 0)
+            label = np.array(label, dtype=np.float32)
+            label = label[np.newaxis, :, :]
+            label[label == 0] = 0
+            label[np.logical_and(label > 0, label < 127.5)] = 2
+            label[label >= 127.5] = 1
+        else:
+            img_file = self.file_list[index].rstrip()
+
+        img = cv2.imread(osp.join(self.root, img_file))
+        img = np.array(img, dtype=np.float32)
+        img = (img - self.mean).transpose((2, 0, 1))
+
+        if self.split == 'train':
+            return img, label
+        else:
+            return img

RCFPyTorch0/models.py

@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import scipy.io as sio
+import torch.nn.functional as F
+
+
+class RCF(nn.Module):
+    def __init__(self, pretrained=None):
+        super(RCF, self).__init__()
+        self.conv1_1 = nn.Conv2d(  3,  64, 3, padding=1, dilation=1)
+        self.conv1_2 = nn.Conv2d( 64,  64, 3, padding=1, dilation=1)
+        self.conv2_1 = nn.Conv2d( 64, 128, 3, padding=1, dilation=1)
+        self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1, dilation=1)
+        self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1, dilation=1)
+        self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1, dilation=1)
+        self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1, dilation=1)
+        self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1, dilation=1)
+        self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1, dilation=1)
+        self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1, dilation=1)
+        self.conv5_1 = nn.Conv2d(512, 512, 3, padding=2, dilation=2)
+        self.conv5_2 = nn.Conv2d(512, 512, 3, padding=2, dilation=2)
+        self.conv5_3 = nn.Conv2d(512, 512, 3, padding=2, dilation=2)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+        self.pool4 = nn.MaxPool2d(2, stride=1, ceil_mode=True)
+        self.act = nn.ReLU(inplace=True)
+
+        self.conv1_1_down = nn.Conv2d( 64, 21, 1)
+        self.conv1_2_down = nn.Conv2d( 64, 21, 1)
+        self.conv2_1_down = nn.Conv2d(128, 21, 1)
+        self.conv2_2_down = nn.Conv2d(128, 21, 1)
+        self.conv3_1_down = nn.Conv2d(256, 21, 1)
+        self.conv3_2_down = nn.Conv2d(256, 21, 1)
+        self.conv3_3_down = nn.Conv2d(256, 21, 1)
+        self.conv4_1_down = nn.Conv2d(512, 21, 1)
+        self.conv4_2_down = nn.Conv2d(512, 21, 1)
+        self.conv4_3_down = nn.Conv2d(512, 21, 1)
+        self.conv5_1_down = nn.Conv2d(512, 21, 1)
+        self.conv5_2_down = nn.Conv2d(512, 21, 1)
+        self.conv5_3_down = nn.Conv2d(512, 21, 1)
+
+        self.score_dsn1 = nn.Conv2d(21, 1, 1)
+        self.score_dsn2 = nn.Conv2d(21, 1, 1)
+        self.score_dsn3 = nn.Conv2d(21, 1, 1)
+        self.score_dsn4 = nn.Conv2d(21, 1, 1)
+        self.score_dsn5 = nn.Conv2d(21, 1, 1)
+        self.score_fuse = nn.Conv2d(5, 1, 1)
+
+        self.weight_deconv2 = self._make_bilinear_weights( 4, 1).cuda()
+        self.weight_deconv3 = self._make_bilinear_weights( 8, 1).cuda()
+        self.weight_deconv4 = self._make_bilinear_weights(16, 1).cuda()
+        self.weight_deconv5 = self._make_bilinear_weights(16, 1).cuda()
+
+        # init weights
+        self.apply(self._init_weights)
+        if pretrained is not None:
+            vgg16 = sio.loadmat(pretrained)
+            torch_params = self.state_dict()
+
+            for k in vgg16.keys():
+                name_par = k.split('-')
+                size = len(name_par)
+                if size == 2:
+                    name_space = name_par[0] + '.' + name_par[1]
+                    data = np.squeeze(vgg16[k])
+                    torch_params[name_space] = torch.from_numpy(data)
+            self.load_state_dict(torch_params)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Conv2d):
+            m.weight.data.normal_(0, 0.01)
+            if m.weight.data.shape == torch.Size([1, 5, 1, 1]):
+                nn.init.constant_(m.weight, 0.2)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    # Based on HED implementation @ https://github.com/xwjabc/hed
+    def _make_bilinear_weights(self, size, num_channels):
+        factor = (size + 1) // 2
+        if size % 2 == 1:
+            center = factor - 1
+        else:
+            center = factor - 0.5
+        og = np.ogrid[:size, :size]
+        filt = (1 - abs(og[0] - center) / factor) * (1 - abs(og[1] - center) / factor)
+        filt = torch.from_numpy(filt)
+        w = torch.zeros(num_channels, num_channels, size, size)
+        w.requires_grad = False
+        for i in range(num_channels):
+            for j in range(num_channels):
+                if i == j:
+                    w[i, j] = filt
+        return w
+
+    # Based on BDCN implementation @ https://github.com/pkuCactus/BDCN
+    def _crop(self, data, img_h, img_w, crop_h, crop_w):
+        _, _, h, w = data.size()
+        assert(img_h <= h and img_w <= w)
+        data = data[:, :, crop_h:crop_h + img_h, crop_w:crop_w + img_w]
+        return data
+
+    def forward(self, x):
+        img_h, img_w = x.shape[2], x.shape[3]
+        conv1_1 = self.act(self.conv1_1(x))
+        conv1_2 = self.act(self.conv1_2(conv1_1))
+        pool1   = self.pool1(conv1_2)
+        conv2_1 = self.act(self.conv2_1(pool1))
+        conv2_2 = self.act(self.conv2_2(conv2_1))
+        pool2   = self.pool2(conv2_2)
+        conv3_1 = self.act(self.conv3_1(pool2))
+        conv3_2 = self.act(self.conv3_2(conv3_1))
+        conv3_3 = self.act(self.conv3_3(conv3_2))
+        pool3   = self.pool3(conv3_3)
+        conv4_1 = self.act(self.conv4_1(pool3))
+        conv4_2 = self.act(self.conv4_2(conv4_1))
+        conv4_3 = self.act(self.conv4_3(conv4_2))
+        pool4   = self.pool4(conv4_3)
+        conv5_1 = self.act(self.conv5_1(pool4))
+        conv5_2 = self.act(self.conv5_2(conv5_1))
+        conv5_3 = self.act(self.conv5_3(conv5_2))
+
+        conv1_1_down = self.conv1_1_down(conv1_1)
+        conv1_2_down = self.conv1_2_down(conv1_2)
+        conv2_1_down = self.conv2_1_down(conv2_1)
+        conv2_2_down = self.conv2_2_down(conv2_2)
+        conv3_1_down = self.conv3_1_down(conv3_1)
+        conv3_2_down = self.conv3_2_down(conv3_2)
+        conv3_3_down = self.conv3_3_down(conv3_3)
+        conv4_1_down = self.conv4_1_down(conv4_1)
+        conv4_2_down = self.conv4_2_down(conv4_2)
+        conv4_3_down = self.conv4_3_down(conv4_3)
+        conv5_1_down = self.conv5_1_down(conv5_1)
+        conv5_2_down = self.conv5_2_down(conv5_2)
+        conv5_3_down = self.conv5_3_down(conv5_3)
+
+        out1 = self.score_dsn1(conv1_1_down + conv1_2_down)
+        out2 = self.score_dsn2(conv2_1_down + conv2_2_down)
+        out3 = self.score_dsn3(conv3_1_down + conv3_2_down + conv3_3_down)
+        out4 = self.score_dsn4(conv4_1_down + conv4_2_down + conv4_3_down)
+        out5 = self.score_dsn5(conv5_1_down + conv5_2_down + conv5_3_down)
+
+        out2 = F.conv_transpose2d(out2, self.weight_deconv2, stride=2)
+        out3 = F.conv_transpose2d(out3, self.weight_deconv3, stride=4)
+        out4 = F.conv_transpose2d(out4, self.weight_deconv4, stride=8)
+        out5 = F.conv_transpose2d(out5, self.weight_deconv5, stride=8)
+
+        out2 = self._crop(out2, img_h, img_w, 1, 1)
+        out3 = self._crop(out3, img_h, img_w, 2, 2)
+        out4 = self._crop(out4, img_h, img_w, 4, 4)
+        out5 = self._crop(out5, img_h, img_w, 0, 0)
+
+        fuse = torch.cat((out1, out2, out3, out4, out5), dim=1)
+        fuse = self.score_fuse(fuse)
+        results = [out1, out2, out3, out4, out5, fuse]
+        results = [torch.sigmoid(r) for r in results]
+        return results

RCFPyTorch0/test.py

@@ -0,0 +1,90 @@
+import os
+import numpy as np
+import os.path as osp
+import cv2
+import argparse
+import torch
+from torch.utils.data import DataLoader
+import torchvision
+from dataset import BSDS_Dataset
+from models import RCF
+
+
+def single_scale_test(model, test_loader, test_list, save_dir):
+    model.eval()
+    if not osp.isdir(save_dir):
+        os.makedirs(save_dir)
+    for idx, image in enumerate(test_loader):
+        image = image.cuda()
+        _, _, H, W = image.shape
+        results = model(image)
+        all_res = torch.zeros((len(results), 1, H, W))
+        for i in range(len(results)):
+          all_res[i, 0, :, :] = results[i]
+        filename = osp.splitext(test_list[idx])[0]
+        torchvision.utils.save_image(1 - all_res, osp.join(save_dir, '%s.jpg' % filename))
+        fuse_res = torch.squeeze(results[-1].detach()).cpu().numpy()
+        fuse_res = ((1 - fuse_res) * 255).astype(np.uint8)
+        cv2.imwrite(osp.join(save_dir, '%s_ss.png' % filename), fuse_res)
+        #print('\rRunning single-scale test [%d/%d]' % (idx + 1, len(test_loader)), end='')
+    print('Running single-scale test done')
+
+
+def multi_scale_test(model, test_loader, test_list, save_dir):
+    model.eval()
+    if not osp.isdir(save_dir):
+        os.makedirs(save_dir)
+    scale = [0.5, 1, 1.5]
+    for idx, image in enumerate(test_loader):
+        in_ = image[0].numpy().transpose((1, 2, 0))
+        _, _, H, W = image.shape
+        ms_fuse = np.zeros((H, W), np.float32)
+        for k in range(len(scale)):
+            im_ = cv2.resize(in_, None, fx=scale[k], fy=scale[k], interpolation=cv2.INTER_LINEAR)
+            im_ = im_.transpose((2, 0, 1))
+            results = model(torch.unsqueeze(torch.from_numpy(im_).cuda(), 0))
+            fuse_res = torch.squeeze(results[-1].detach()).cpu().numpy()
+            fuse_res = cv2.resize(fuse_res, (W, H), interpolation=cv2.INTER_LINEAR)
+            ms_fuse += fuse_res
+        ms_fuse = ms_fuse / len(scale)
+        ### rescale trick
+        # ms_fuse = (ms_fuse - ms_fuse.min()) / (ms_fuse.max() - ms_fuse.min())
+        filename = osp.splitext(test_list[idx])[0]
+        ms_fuse = ((1 - ms_fuse) * 255).astype(np.uint8)
+        cv2.imwrite(osp.join(save_dir, '%s_ms.png' % filename), ms_fuse)
+        #print('\rRunning multi-scale test [%d/%d]' % (idx + 1, len(test_loader)), end='')
+    print('Running multi-scale test done')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='PyTorch Testing')
+    parser.add_argument('--gpu', default='0', type=str, help='GPU ID')
+    parser.add_argument('--checkpoint', default=None, type=str, help='path to latest checkpoint')
+    parser.add_argument('--save-dir', help='output folder', default='results/RCF')
+    parser.add_argument('--dataset', help='root folder of dataset', default='data/HED-BSDS')
+    args = parser.parse_args()
+
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+
+    if not osp.isdir(args.save_dir):
+        os.makedirs(args.save_dir)
+  
+    test_dataset  = BSDS_Dataset(root=args.dataset, split='test')
+    test_loader   = DataLoader(test_dataset, batch_size=1, num_workers=1, drop_last=False, shuffle=False)
+    test_list = [osp.split(i.rstrip())[1] for i in test_dataset.file_list]
+    assert len(test_list) == len(test_loader)
+
+    model = RCF().cuda()
+
+    if osp.isfile(args.checkpoint):
+        print("=> loading checkpoint from '{}'".format(args.checkpoint))
+        checkpoint = torch.load(args.checkpoint)
+        model.load_state_dict(checkpoint)
+        print("=> checkpoint loaded")
+    else:
+        print("=> no checkpoint found at '{}'".format(args.checkpoint))
+
+    print('Performing the testing...')
+    single_scale_test(model, test_loader, test_list, args.save_dir)
+    multi_scale_test(model, test_loader, test_list, args.save_dir)

RCFPyTorch0/train.py

@@ -0,0 +1,209 @@
+import os
+import numpy as np
+import os.path as osp
+import cv2
+import argparse
+import time
+import torch
+from torch.utils.data import DataLoader
+import torchvision
+from dataset import BSDS_Dataset
+from models import RCF
+from utils import Logger, Averagvalue, Cross_entropy_loss
+
+
+def train(args, model, train_loader, optimizer, epoch, logger):
+    batch_time = Averagvalue()
+    losses = Averagvalue()
+    model.train()
+    end = time.time()
+    counter = 0
+    for i, (image, label) in enumerate(train_loader):
+        image, label = image.cuda(), label.cuda()
+        outputs = model(image)
+        loss = torch.zeros(1).cuda()
+        for o in outputs:
+            loss = loss + Cross_entropy_loss(o, label)
+        counter += 1
+        loss = loss / args.iter_size
+        loss.backward()
+        if counter == args.iter_size:
+            optimizer.step()
+            optimizer.zero_grad()
+            counter = 0
+        # measure accuracy and record loss
+        losses.update(loss.item(), image.size(0))
+        batch_time.update(time.time() - end)
+        if i % args.print_freq == 0:
+            logger.info('Epoch: [{0}/{1}][{2}/{3}] '.format(epoch + 1, args.max_epoch, i, len(train_loader)) + \
+                        'Time {batch_time.val:.3f} (avg: {batch_time.avg:.3f}) '.format(batch_time=batch_time) + \
+                        'Loss {loss.val:f} (avg: {loss.avg:f}) '.format(loss=losses))
+        end = time.time()
+
+
+def single_scale_test(model, test_loader, test_list, save_dir):
+    model.eval()
+    if not osp.isdir(save_dir):
+        os.makedirs(save_dir)
+    for idx, image in enumerate(test_loader):
+        image = image.cuda()
+        _, _, H, W = image.shape
+        results = model(image)
+        all_res = torch.zeros((len(results), 1, H, W))
+        for i in range(len(results)):
+          all_res[i, 0, :, :] = results[i]
+        filename = osp.splitext(test_list[idx])[0]
+        torchvision.utils.save_image(1 - all_res, osp.join(save_dir, '%s.jpg' % filename))
+        fuse_res = torch.squeeze(results[-1].detach()).cpu().numpy()
+        fuse_res = ((1 - fuse_res) * 255).astype(np.uint8)
+        cv2.imwrite(osp.join(save_dir, '%s_ss.png' % filename), fuse_res)
+        #print('\rRunning single-scale test [%d/%d]' % (idx + 1, len(test_loader)), end='')
+    logger.info('Running single-scale test done')
+
+
+def multi_scale_test(model, test_loader, test_list, save_dir):
+    model.eval()
+    if not osp.isdir(save_dir):
+        os.makedirs(save_dir)
+    scale = [0.5, 1, 1.5]
+    for idx, image in enumerate(test_loader):
+        in_ = image[0].numpy().transpose((1, 2, 0))
+        _, _, H, W = image.shape
+        ms_fuse = np.zeros((H, W), np.float32)
+        for k in range(len(scale)):
+            im_ = cv2.resize(in_, None, fx=scale[k], fy=scale[k], interpolation=cv2.INTER_LINEAR)
+            im_ = im_.transpose((2, 0, 1))
+            results = model(torch.unsqueeze(torch.from_numpy(im_).cuda(), 0))
+            fuse_res = torch.squeeze(results[-1].detach()).cpu().numpy()
+            fuse_res = cv2.resize(fuse_res, (W, H), interpolation=cv2.INTER_LINEAR)
+            ms_fuse += fuse_res
+        ms_fuse = ms_fuse / len(scale)
+        ### rescale trick
+        # ms_fuse = (ms_fuse - ms_fuse.min()) / (ms_fuse.max() - ms_fuse.min())
+        filename = osp.splitext(test_list[idx])[0]
+        ms_fuse = ((1 - ms_fuse) * 255).astype(np.uint8)
+        cv2.imwrite(osp.join(save_dir, '%s_ms.png' % filename), ms_fuse)
+        #print('\rRunning multi-scale test [%d/%d]' % (idx + 1, len(test_loader)), end='')
+    logger.info('Running multi-scale test done')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='PyTorch Training')
+    parser.add_argument('--batch-size', default=1, type=int, help='batch size')
+    parser.add_argument('--lr', default=1e-6, type=float, help='initial learning rate')
+    parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
+    parser.add_argument('--weight-decay', default=2e-4, type=float, help='weight decay')
+    parser.add_argument('--stepsize', default=3, type=int, help='learning rate step size')
+    parser.add_argument('--gamma', default=0.1, type=float, help='learning rate decay rate')
+    parser.add_argument('--max-epoch', default=10, type=int, help='the number of training epochs')
+    parser.add_argument('--iter-size', default=10, type=int, help='iter size')
+    parser.add_argument('--start-epoch', default=0, type=int, help='manual epoch number')
+    parser.add_argument('--print-freq', default=200, type=int, help='print frequency')
+    parser.add_argument('--gpu', default='0', type=str, help='GPU ID')
+    parser.add_argument('--resume', default=None, type=str, help='path to latest checkpoint')
+    parser.add_argument('--save-dir', help='output folder', default='results/RCF')
+    parser.add_argument('--dataset', help='root folder of dataset', default='data')
+    args = parser.parse_args()
+
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+
+    if not osp.isdir(args.save_dir):
+        os.makedirs(args.save_dir)
+
+    logger = Logger(osp.join(args.save_dir, 'log.txt'))
+    logger.info('Called with args:')
+    for (key, value) in vars(args).items():
+        logger.info('{0:15} | {1}'.format(key, value))
+
+    train_dataset = BSDS_Dataset(root=args.dataset, split='train')
+    test_dataset  = BSDS_Dataset(root=osp.join(args.dataset, 'HED-BSDS'), split='test')
+    train_loader  = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=4, drop_last=True, shuffle=True)
+    test_loader   = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=4, drop_last=False, shuffle=False)
+    test_list = [osp.split(i.rstrip())[1] for i in test_dataset.file_list]
+    assert len(test_list) == len(test_loader)
+
+    model = RCF(pretrained='vgg16convs.mat').cuda()
+    parameters = {'conv1-4.weight': [], 'conv1-4.bias': [], 'conv5.weight': [], 'conv5.bias': [],
+        'conv_down_1-5.weight': [], 'conv_down_1-5.bias': [], 'score_dsn_1-5.weight': [],
+        'score_dsn_1-5.bias': [], 'score_fuse.weight': [], 'score_fuse.bias': []}
+    for pname, p in model.named_parameters():
+        if pname in ['conv1_1.weight','conv1_2.weight',
+                     'conv2_1.weight','conv2_2.weight',
+                     'conv3_1.weight','conv3_2.weight','conv3_3.weight',
+                     'conv4_1.weight','conv4_2.weight','conv4_3.weight']:
+            parameters['conv1-4.weight'].append(p)
+        elif pname in ['conv1_1.bias','conv1_2.bias',
+                       'conv2_1.bias','conv2_2.bias',
+                       'conv3_1.bias','conv3_2.bias','conv3_3.bias',
+                       'conv4_1.bias','conv4_2.bias','conv4_3.bias']:
+            parameters['conv1-4.bias'].append(p)
+        elif pname in ['conv5_1.weight','conv5_2.weight','conv5_3.weight']:
+            parameters['conv5.weight'].append(p)
+        elif pname in ['conv5_1.bias','conv5_2.bias','conv5_3.bias']:
+            parameters['conv5.bias'].append(p)
+        elif pname in ['conv1_1_down.weight','conv1_2_down.weight',
+                       'conv2_1_down.weight','conv2_2_down.weight',
+                       'conv3_1_down.weight','conv3_2_down.weight','conv3_3_down.weight',
+                       'conv4_1_down.weight','conv4_2_down.weight','conv4_3_down.weight',
+                       'conv5_1_down.weight','conv5_2_down.weight','conv5_3_down.weight']:
+            parameters['conv_down_1-5.weight'].append(p)
+        elif pname in ['conv1_1_down.bias','conv1_2_down.bias',
+                       'conv2_1_down.bias','conv2_2_down.bias',
+                       'conv3_1_down.bias','conv3_2_down.bias','conv3_3_down.bias',
+                       'conv4_1_down.bias','conv4_2_down.bias','conv4_3_down.bias',
+                       'conv5_1_down.bias','conv5_2_down.bias','conv5_3_down.bias']:
+            parameters['conv_down_1-5.bias'].append(p)
+        elif pname in ['score_dsn1.weight','score_dsn2.weight','score_dsn3.weight', 'score_dsn4.weight','score_dsn5.weight']:
+            parameters['score_dsn_1-5.weight'].append(p)
+        elif pname in ['score_dsn1.bias','score_dsn2.bias','score_dsn3.bias', 'score_dsn4.bias','score_dsn5.bias']:
+            parameters['score_dsn_1-5.bias'].append(p)
+        elif pname in ['score_fuse.weight']:
+            parameters['score_fuse.weight'].append(p)
+        elif pname in ['score_fuse.bias']:
+            parameters['score_fuse.bias'].append(p)
+
+    optimizer = torch.optim.SGD([
+            {'params': parameters['conv1-4.weight'],       'lr': args.lr*1,     'weight_decay': args.weight_decay},
+            {'params': parameters['conv1-4.bias'],         'lr': args.lr*2,     'weight_decay': 0.},
+            {'params': parameters['conv5.weight'],         'lr': args.lr*100,   'weight_decay': args.weight_decay},
+            {'params': parameters['conv5.bias'],           'lr': args.lr*200,   'weight_decay': 0.},
+            {'params': parameters['conv_down_1-5.weight'], 'lr': args.lr*0.1,   'weight_decay': args.weight_decay},
+            {'params': parameters['conv_down_1-5.bias'],   'lr': args.lr*0.2,   'weight_decay': 0.},
+            {'params': parameters['score_dsn_1-5.weight'], 'lr': args.lr*0.01,  'weight_decay': args.weight_decay},
+            {'params': parameters['score_dsn_1-5.bias'],   'lr': args.lr*0.02,  'weight_decay': 0.},
+            {'params': parameters['score_fuse.weight'],    'lr': args.lr*0.001, 'weight_decay': args.weight_decay},
+            {'params': parameters['score_fuse.bias'],      'lr': args.lr*0.002, 'weight_decay': 0.},
+        ], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)
+
+    if args.resume is not None:
+        if osp.isfile(args.resume):
+            logger.info("=> loading checkpoint from '{}'".format(args.resume))
+            checkpoint = torch.load(args.resume)
+            model.load_state_dict(checkpoint['state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            logger.info("=> checkpoint loaded")
+        else:
+            logger.info("=> no checkpoint found at '{}'".format(args.resume))
+
+    for epoch in range(args.start_epoch, args.max_epoch):
+        logger.info('Performing initial testing...')
+        train(args, model, train_loader, optimizer, epoch, logger)
+        save_dir = osp.join(args.save_dir, 'epoch%d-test' % (epoch + 1))
+        single_scale_test(model, test_loader, test_list, save_dir)
+        multi_scale_test(model, test_loader, test_list, save_dir)
+        # Save checkpoint
+        save_file = osp.join(args.save_dir, 'checkpoint_epoch{}.pth'.format(epoch + 1))
+        torch.save({
+                'epoch': epoch,
+                'args': args,
+                'state_dict': model.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'lr_scheduler': lr_scheduler.state_dict(),
+            }, save_file)
+        lr_scheduler.step() # will adjust learning rate
+
+    logger.close()

RCFPyTorch0/utils.py

@@ -0,0 +1,54 @@
+import os
+import logging
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class Logger(object):
+    def __init__(self, path='log.txt'):
+        self.logger = logging.getLogger('Logger')
+        self.file_handler = logging.FileHandler(path, 'w')
+        self.stdout_handler = logging.StreamHandler()
+        self.logger.addHandler(self.file_handler)
+        self.logger.addHandler(self.stdout_handler)
+        self.stdout_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+        self.file_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+        self.logger.setLevel(logging.INFO)
+
+    def info(self, txt):
+        self.logger.info(txt)
+
+    def close(self):
+        self.file_handler.close()
+        self.stdout_handler.close()
+
+
+class Averagvalue(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def Cross_entropy_loss(prediction, label):
+    mask = label.clone()
+    num_positive = torch.sum((mask == 1).float()).float()
+    num_negative = torch.sum((mask == 0).float()).float()
+
+    mask[mask == 1] = 1.0 * num_negative / (num_positive + num_negative)
+    mask[mask == 0] = 1.1 * num_positive / (num_positive + num_negative)
+    mask[mask == 2] = 0
+    cost = F.binary_cross_entropy(prediction, label, weight=mask, reduce=False)
+    return torch.sum(cost)

RCFPyTorch0/vgg16convs.mat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bce56b30c32d4c72954355fe970c87dceba15bc180aa89524960fda1e0e32cd9
+size 58860856