spikezoo 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

spikezoo/archs/spk2imgnet/dataset.py

@@ -0,0 +1,144 @@
+import os
+import os.path
+from typing import List
+
+import numpy as np
+import random
+import h5py
+import torch
+import cv2
+import glob
+import torch.utils.data as udata
+from functools import partial
+
+bytes2num = partial(int.from_bytes, byteorder="little", signed=False)
+
+
+def normalize(data):
+    return data / 255.0
+
+
+def raw_to_spike(video_seq, h, w):
+    video_seq = np.array(video_seq).astype(np.uint8)
+    img_size = h * w
+    img_num = len(video_seq) // (img_size // 8)
+    spike_matrix = np.zeros([img_num, h, w], np.uint8)
+    pix_id = np.arange(0, h * w)
+    pix_id = np.reshape(pix_id, (h, w))
+    comparator = np.left_shift(1, np.mod(pix_id, 8))
+    byte_id = pix_id // 8
+
+    for img_id in np.arange(img_num):
+        id_start = img_id * img_size // 8
+        id_end = id_start + img_size // 8
+        cur_info = video_seq[id_start:id_end]
+        data = cur_info[byte_id]
+        result = np.bitwise_and(data, comparator)
+        spike_matrix[img_id, :, :] = np.flipud((result == comparator))
+
+    return spike_matrix
+
+
+def Im2Patch(img, win, stride=40):
+    k = 0
+    [endc, endw, endh] = img.shape
+    patch = img[:, 0: endw - win + 0 + 1: stride, 0: endh - win + 0 + 1: stride]
+    total_pat_num = patch.shape[1] * patch.shape[2]
+    Y = np.zeros([endc, win * win, total_pat_num], np.float32)
+    for i in range(win):
+        for j in range(win):
+            patch = img[
+                    :, i: endw - win + i + 1: stride, j: endh - win + j + 1: stride
+                    ]
+            Y[:, k, :] = np.array(patch[:]).reshape(endc, total_pat_num)
+            k = k + 1
+    return Y.reshape([endc, win, win, total_pat_num])
+
+
+def read_image_and_concat_as_tensor(paths: List[str]):
+    tensors = []
+    for path in paths:
+        img = cv2.imread(path)
+        tensors.append(img.reshape([1, *img.shape]))
+    return np.concatenate(tensors, axis=0)
+
+
+def prepare_data(data_path, patch_size, stride, h5_name, aug_times=1):
+    print("process training data")
+    input_files = glob.glob(os.path.join(data_path, "input", "*.dat"))
+    print(len(input_files))
+    input_files.sort()
+    input_h5f = h5py.File(h5_name + "_input.h5", "w")
+    gt_h5f = h5py.File(h5_name + "_gt.h5", "w")
+    train_num = 0
+    h = 250
+    w = 400
+    for i in range(len(input_files)):
+        input_f = open(input_files[i], "rb+")
+        video_seq = input_f.read()
+        video_seq = np.fromstring(video_seq, "B")
+        # print(video_seq)
+        spike_array = raw_to_spike(video_seq, h, w)  # c*h*w
+        # print(input_files[i][:-3])
+        # SpikeArray = SpikeArray[10:-10, :, :]
+        # print(np.mean(SpikeArray))
+        print(spike_array.shape)
+        file_name = input_files[i].replace("\\", "/").split("/")[-1]
+        gt = []
+        for num in [7, 14, 21, 28, 35]:
+            img = cv2.imread(os.path.join(data_path, "gt", file_name[:-6] + str(num) + ".png"), 0)
+            gt.append(img.reshape([1, *img.shape]))
+        gt = np.concatenate(gt, axis=0)
+        print(input_files[i])
+        print(os.path.join(data_path, "gt", file_name[:-3] + "png"))
+        gt = np.float32(normalize(gt))  # size
+        print(gt.shape)
+        print(spike_array.shape)
+        input_patches = Im2Patch(spike_array, win=patch_size, stride=stride)
+        gt_patches = Im2Patch(gt, win=patch_size, stride=stride)
+        assert input_patches.shape[3] == gt_patches.shape[3]
+        for n in range(input_patches.shape[3]):
+            inputs = input_patches[:, :, :, n].copy()
+            input_h5f.create_dataset(str(train_num), data=inputs)
+            gt = gt_patches[:, :, :, n].copy()
+            gt_h5f.create_dataset(str(train_num), data=gt)
+            train_num += 1
+
+    input_h5f.close()
+    gt_h5f.close()
+
+
+class Dataset(udata.Dataset):
+    def __init__(self, h5_name):
+        super(Dataset, self).__init__()
+        input_h5f = h5py.File(h5_name + "_input.h5", "r")
+        gt_h5f = h5py.File(h5_name + "_gt.h5", "r")
+        self.h5_name = h5_name
+        self.keys = list(input_h5f.keys())
+        # print(self.keys)
+        random.shuffle(self.keys)
+        input_h5f.close()
+        gt_h5f.close()
+
+    def __len__(self):
+        return len(self.keys)
+
+    def __getitem__(self, index):
+        input_h5f = h5py.File(self.h5_name + "_input.h5", "r")
+        gt_h5f = h5py.File(self.h5_name + "_gt.h5", "r")
+        key = self.keys[index]
+        inputs = np.array(input_h5f[key])
+        gt = np.array(gt_h5f[key])
+        input_h5f.close()
+        gt_h5f.close()
+        return torch.Tensor(inputs), torch.Tensor(gt)
+
+
+if __name__ == "__main__":
+    prepare_data(
+        data_path="./Spk2ImgNet_train/train2/",
+        patch_size=40,
+        stride=40,
+        h5_name="train",
+    )
+    # PrepareData(data_path = './SpikeDataset/val/', patch_size=40, stride=40, h5_name='val')

spikezoo/archs/spk2imgnet/nets.py

@@ -0,0 +1,230 @@
+import numpy as np
+import sys
+import os
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(current_dir)
+
+from align_arch import *
+
+class BasicBlock(nn.Module):
+    def __init__(self, features):
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels=features,
+            out_channels=features,
+            kernel_size=3,
+            padding=1,
+            bias=True,
+        )
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv2d(
+            in_channels=features,
+            out_channels=features,
+            kernel_size=3,
+            padding=1,
+            bias=True,
+        )
+        self.relu2 = nn.ReLU()
+        self.conv3 = nn.Conv2d(
+            in_channels=features,
+            out_channels=features,
+            kernel_size=3,
+            padding=1,
+            bias=True,
+        )
+        self.relu3 = nn.ReLU()
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.relu1(out)
+        out = self.conv2(out)
+        out = self.relu2(out)
+        out = self.conv3(out)
+        return self.relu3(x + out)
+
+
+# use Sigmoid
+class CALayer2(nn.Module):
+    def __init__(self, in_channels):
+        super(CALayer2, self).__init__()
+        self.ca_block = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels * 2, 3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels * 2, in_channels, 3, padding=1, bias=True),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        weight = self.ca_block(x)
+        return weight
+
+
+# use CALayer
+class FeatureExtractor(nn.Module):
+    def __init__(
+        self, in_channels, features, out_channels, channel_step, num_of_layers=16
+    ):
+        super(FeatureExtractor, self).__init__()
+        # self.InferLayer = LightInferLayer(in_channels=in_channels)
+        self.channel_step = channel_step
+        self.conv0_0 = nn.Conv2d(
+            in_channels=in_channels, out_channels=16, kernel_size=3, padding=1
+        )
+        self.conv0_1 = nn.Conv2d(
+            in_channels=in_channels - 2 * channel_step,
+            out_channels=16,
+            kernel_size=3,
+            padding=1,
+        )
+        self.conv0_2 = nn.Conv2d(
+            in_channels=in_channels - 4 * channel_step,
+            out_channels=16,
+            kernel_size=3,
+            padding=1,
+        )
+        self.conv0_3 = nn.Conv2d(
+            in_channels=in_channels - 6 * channel_step,
+            out_channels=16,
+            kernel_size=3,
+            padding=1,
+        )
+        self.conv1_0 = nn.Conv2d(
+            in_channels=16, out_channels=1, kernel_size=3, padding=1
+        )
+        self.conv1_1 = nn.Conv2d(
+            in_channels=16, out_channels=1, kernel_size=3, padding=1
+        )
+        self.conv1_2 = nn.Conv2d(
+            in_channels=16, out_channels=1, kernel_size=3, padding=1
+        )
+        self.conv1_3 = nn.Conv2d(
+            in_channels=16, out_channels=1, kernel_size=3, padding=1
+        )
+        self.ca = CALayer2(in_channels=4)
+        self.conv = nn.Conv2d(
+            in_channels=4, out_channels=features, kernel_size=3, padding=1
+        )
+        self.relu = nn.ReLU()
+        layers = []
+        for _ in range(num_of_layers - 2):
+            layers.append(BasicBlock(features=features))
+        # layers.append(nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=kernel_size, padding=padding, bias=True))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        out_0 = self.conv1_0(self.relu(self.conv0_0(x)))
+        out_1 = self.conv1_1(
+            self.relu(self.conv0_1(x[:, self.channel_step : -self.channel_step, :, :]))
+        )
+        out_2 = self.conv1_2(
+            self.relu(
+                self.conv0_2(x[:, 2 * self.channel_step : -2 * self.channel_step, :, :])
+            )
+        )
+        out_3 = self.conv1_3(
+            self.relu(
+                self.conv0_3(x[:, 3 * self.channel_step : -3 * self.channel_step, :, :])
+            )
+        )
+        out = torch.cat((out_0, out_1), 1)
+        out = torch.cat((out, out_2), 1)
+        out = torch.cat((out, out_3), 1)
+        est = out
+        weight = self.ca(out)
+        out = weight * out
+        out = self.conv(out)
+        out = self.relu(out)
+        tmp = out
+        out = self.net(out)
+        # out = self.conv2(out)
+        # out = self.relu2(out)
+        # out = self.conv3(out)
+        return out + tmp, est
+
+
+class FusionMaskV1(nn.Module):
+    def __init__(self, features):
+        super(FusionMaskV1, self).__init__()
+        self.conv0 = nn.Conv2d(
+            in_channels=2 * features, out_channels=features, kernel_size=3, padding=1
+        )
+        self.conv1 = nn.Conv2d(
+            in_channels=features, out_channels=features, kernel_size=3, padding=1
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=features, out_channels=features, kernel_size=3, padding=1
+        )
+        self.prelu0 = nn.PReLU()
+        self.prelu1 = nn.PReLU()
+        self.sig = nn.Sigmoid()
+
+    def forward(self, ref, key):
+        fea = torch.cat((ref, key), 1)
+        fea = self.conv2(self.prelu1(self.conv1(self.prelu0(self.conv0(fea)))))
+        mask = self.sig(fea)
+        return mask
+
+
+# current best model
+class SpikeNet(nn.Module):
+    def __init__(self, in_channels, features, out_channels, win_r, win_step):
+        super(SpikeNet, self).__init__()
+        self.extractor = FeatureExtractor(
+            in_channels=in_channels,
+            features=features,
+            out_channels=features,
+            channel_step=1,
+            num_of_layers=12,
+        )
+        self.mask0 = FusionMaskV1(features=features)
+        self.mask1 = FusionMaskV1(features=features)
+        self.mask3 = FusionMaskV1(features=features)
+        self.mask4 = FusionMaskV1(features=features)
+        self.rec_conv0 = nn.Conv2d(
+            in_channels=5 * features,
+            out_channels=3 * features,
+            kernel_size=3,
+            padding=1,
+        )
+        self.rec_conv1 = nn.Conv2d(
+            in_channels=3 * features, out_channels=features, kernel_size=3, padding=1
+        )
+        self.rec_conv2 = nn.Conv2d(
+            in_channels=features, out_channels=1, kernel_size=3, padding=1
+        )
+        self.rec_relu = nn.ReLU()
+        self.pcd_align = Easy_PCD(nf=features, groups=8)
+        self.win_r = win_r
+        self.win_step = win_step
+
+    def forward(self, x):
+        block0 = x[:, 0 : 2 * self.win_r + 1, :, :]
+        block1 = x[:, self.win_step : self.win_step + 2 * self.win_r + 1, :, :]
+        block2 = x[:, 2 * self.win_step : 2 * self.win_step + 2 * self.win_r + 1, :, :]
+        block3 = x[:, 3 * self.win_step : 3 * self.win_step + 2 * self.win_r + 1, :, :]
+        block4 = x[:, 4 * self.win_step : 4 * self.win_step + 2 * self.win_r + 1, :, :]
+        block0_out, est0 = self.extractor(block0)
+        block1_out, est1 = self.extractor(block1)
+        block2_out, est2 = self.extractor(block2)
+        block3_out, est3 = self.extractor(block3)
+        block4_out, est4 = self.extractor(block4)
+        aligned_block0_out = self.pcd_align(block0_out, block2_out)
+        aligned_block1_out = self.pcd_align(block1_out, block2_out)
+        aligned_block3_out = self.pcd_align(block3_out, block2_out)
+        aligned_block4_out = self.pcd_align(block4_out, block2_out)
+        mask0 = self.mask0(aligned_block0_out, block2_out)
+        mask1 = self.mask1(aligned_block1_out, block2_out)
+        mask3 = self.mask3(aligned_block3_out, block2_out)
+        mask4 = self.mask4(aligned_block4_out, block2_out)
+        out = torch.cat((aligned_block0_out * mask0, aligned_block1_out * mask1), 1)
+        out = torch.cat((out, block2_out), 1)
+        out = torch.cat((out, aligned_block3_out * mask3), 1)
+        out = torch.cat((out, aligned_block4_out * mask4), 1)
+        out = self.rec_relu(self.rec_conv0(out))
+        out = self.rec_relu(self.rec_conv1(out))
+        out = self.rec_conv2(out)
+        return out
+
+
+if __name__ == "__main__":
+    print("out")

spikezoo/archs/spk2imgnet/readme.md

@@ -0,0 +1,86 @@
+## [CVPR 2021] Spk2ImgNet: Learning to Reconstruct Dynamic Scene from Continuous Spike Stream
+
+
+<h4 align="center"> Jing Zhao, Ruiqin Xiong, Hangfan Liu, Jian Zhang, Tiejun Huang </h4>
+
+This repository contains the official source code for our paper:
+
+Spk2ImgNet: Learning to Reconstruct Dynamic Scene from Continuous Spike Stream.  CVPR 2021 
+
+Paper:  
+[Spk2ImgNet-CVPR2021](https://openaccess.thecvf.com/content/CVPR2021/papers/Zhao_Spk2ImgNet_Learning_To_Reconstruct_Dynamic_Scene_From_Continuous_Spike_Stream_CVPR_2021_paper.pdf) 
+
+* [Spk2ImgNet](#Learning-to-Reconstruct-Dynamic-Scene-from-Continuous-Spike-Stream.)
+  * [Environments](#Environments)
+  * [Download the pretrained models](#Download-the-pretrained-models)
+  * [Evaluate](#Evaluate)
+  * [Train](#Train)
+  * [Citation](#Citations)
+
+
+## Environments
+
+You will have to choose cudatoolkit version to match your compute environment. The code is tested on PyTorch 1.10.2+cu113 and spatial-correlation-sampler 0.3.0 but other versions might also work. 
+
+```bash
+conda create -n steflow python==3.9
+conda activate steflow
+conda install pytorch torchvision cudatoolkit=11.3 -c pytorch
+pip3 install matplotlib opencv-python h5py
+```
+
+We don't ensure that all the PyTorch versions can work well.
+
+## Prepare the Data
+
+### Download the pretrained models
+
+The pretrained model can be downloaded in the Google Drive link below
+
+[Link for pretrained model](https://drive.google.com/file/d/1vBTJxlctk4otQKsyRq7lsFYGU4WGRNjt/view?usp=sharing)
+
+You can download the pretrained models to ```./ckpt```
+
+### Download the training data
+
+The training data can be downloaded in the Google Drive link below
+
+[Link for training data](https://drive.google.com/file/d/1ozR2-fNmU10gA_TCYUfJN-ahV6e_8Ke7/view?usp=sharing)
+
+## Evaluate 
+
+You can set the data path in the .py files or through argparser (--data)
+
+```bash
+python3 main_steflow_dt1.py \
+--test_data 'Spk2ImgNet_test2' \
+--model_name 'model_061.pth'
+
+```
+
+
+## Train
+
+
+All the command line arguments for hyperparameter tuning can be found in the `train.py` file.
+You can set the data path in the .py files or through argparser (--data)
+
+```bash
+python3 train.py
+```
+
+## Citations
+
+If you find this code useful in your research, please consider citing our paper: 
+
+```
+@inproceedings{zhao2021spike,
+  title={Spk2ImgNet: Learning to Reconstruct Dynamic Scene from Continuous Spike Stream},
+  author={Zhao, Jing and Xiong, Ruiqin and Liu, Hangfan and Zhang, Jian and Huang, Tiejun},
+  booktitle={2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2021}
+}
+```
+
+
+

spikezoo/archs/spk2imgnet/test_gen_imgseq.py

@@ -0,0 +1,118 @@
+import argparse
+import time
+
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+from torch.autograd import Variable
+
+from dataset import *
+from nets import *
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+parser = argparse.ArgumentParser(description="Spike_Net_Test")
+parser.add_argument(
+    "--num_of_layers", type=int, default=17, help="Number of toatal layers"
+)
+parser.add_argument(
+    "--logdir",
+    type=str,
+    default="./ckpt2/",
+    help="path of log files",
+)
+parser.add_argument("--test_data", type=str, default="./Spk2ImgNet_test2/test2/", help="test set")
+parser.add_argument(
+    "--save_result", type=bool, default=True, help="save the reconstruction or not"
+)
+parser.add_argument(
+    "--result_dir", type=str, default="results/", help="path of results"
+)
+parser.add_argument(
+    "--exist_gt", type=bool, default=True, help="exist ground truth or not"
+)
+parser.add_argument("--model_name", type=str, default="model_041.pth", help="Name of ckp")
+opt = parser.parse_args()
+
+
+def normalize(data):
+    return data / 255.0
+
+
+def main():
+    # Build model
+    print("Loading model ... \n")
+    net = SpikeNet(
+        in_channels=13, features=64, out_channels=1, win_r=6, win_step=7
+    )
+    # device_ids = [0]
+    # print(device_ids[0])
+    model = nn.DataParallel(net).cuda()
+    model.load_state_dict(torch.load(os.path.join(opt.logdir, opt.model_name)))
+    model.eval()
+
+    # load data info
+    print("Loading data info ...\n")
+    # sub_dir = 'data4'
+    files_source = glob.glob(os.path.join(opt.test_data, "input", "*.dat"))
+    files_source.sort()
+
+    # process data
+    psnr_test = 0
+    ssim_test = 0
+    for i in range(len(files_source)):
+        sub_dir = files_source[i][:-4]
+        # Input spike
+        input_f = open(files_source[i], "rb+")
+        video_seq = input_f.read()
+        video_seq = np.fromstring(video_seq, "B")
+        InSpikeArray = raw_to_spike(video_seq, 250, 400)  # c*h*w
+        [c, h, w] = InSpikeArray.shape
+        for key_id in np.arange(151, 152, 1):
+            start_t = time.time()
+            SpikeArray = InSpikeArray[key_id - 21 : key_id + 20, :, :]
+            # make its shape can be divided by 4
+            SpikeArray = np.pad(
+                SpikeArray, ((0, 0), (0, 2), (0, 0)), "symmetric"
+            )  # c*252*40
+            SpikeArray = np.expand_dims(SpikeArray, 0)  # n*c*h*w
+            file_name = files_source[i].replace("\\", "/").split("/")[-1]
+
+            SpikeArray = Variable(torch.Tensor(SpikeArray)).cuda()
+            with torch.no_grad():
+                if opt.exist_gt:
+                    out_rec, est0, est1, est2, est3, est4 = model(SpikeArray)
+                    out_rec = (
+                        torch.clamp(out_rec / 0.6, 0, 1).cpu() * 255
+                    )  # 0.6 is the converation rate used in the spike camera. Only neccessary for our synthezed data.
+                else:
+                    out_rec, est0, est1, est2, est3, est4 = model(SpikeArray)
+                    out_rec = torch.clamp(out_rec, 0, 1).cpu() ** (1 / 2.2) * 255
+            out_rec = out_rec.detach().numpy().astype(np.float32)
+            out_rec = np.squeeze(out_rec).astype(np.uint8)
+            # transform to orignal shape # 250*400
+            out_rec = out_rec[:250, :]
+            if opt.exist_gt:
+                gt = cv2.imread(
+                    os.path.join(opt.test_data, "gt", file_name[:-3] + "png"), 0
+                )
+                psnr = peak_signal_noise_ratio(gt, out_rec)
+                ssim = structural_similarity(gt, out_rec)
+                print("%10s: PSNR:%.2f SSIM:%.4f" % (file_name, psnr, ssim))
+                psnr_test += psnr
+                ssim_test += ssim
+            if opt.save_result:
+                if not os.path.exists(os.path.join(opt.result_dir, sub_dir)):
+                    os.makedirs(os.path.join(opt.result_dir, sub_dir))
+                cv2.imwrite(
+                    os.path.join(opt.result_dir, sub_dir, str(key_id) + ".png"), out_rec
+                )
+                dur_time = time.time() - start_t
+                print("dur_time:%.2f", dur_time)
+
+    if opt.exist_gt:
+        avg_psnr = psnr_test / len(files_source)
+        avg_ssim = ssim_test / len(files_source)
+        print("average PSNR: %.2f average SSIM: %.4f" % (avg_psnr, avg_ssim))
+
+
+if __name__ == "__main__":
+    main()