spikezoo 0.1.1__py3-none-any.whl → 0.2__py3-none-any.whl

spikezoo/archs/stir/package_core/package_core/disp_netS.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, zeros_
+
+
+def downsample_conv(in_planes, out_planes, kernel_size=3):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=2, padding=(kernel_size-1)//2),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(out_planes, out_planes, kernel_size=kernel_size, padding=(kernel_size-1)//2),
+        nn.ReLU(inplace=True)
+    )
+
+
+def predict_disp(in_planes):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, 1, kernel_size=3, padding=1),
+        nn.Sigmoid()
+    )
+
+
+def conv(in_planes, out_planes):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=3, padding=1),
+        nn.ReLU(inplace=True)
+    )
+
+
+def upconv(in_planes, out_planes):
+    return nn.Sequential(
+        nn.ConvTranspose2d(in_planes, out_planes, kernel_size=3, stride=2, padding=1, output_padding=1),
+        nn.ReLU(inplace=True)
+    )
+
+
+def crop_like(input, ref):
+    assert(input.size(2) >= ref.size(2) and input.size(3) >= ref.size(3))
+    return input[:, :, :ref.size(2), :ref.size(3)]
+
+
+class DispNetS(nn.Module):
+
+    def __init__(self, alpha=10, beta=0.01):
+        super(DispNetS, self).__init__()
+
+        self.alpha = alpha
+        self.beta = beta
+
+        conv_planes = [32, 64, 128, 256, 512, 512, 512]
+        self.conv1 = downsample_conv(3,              conv_planes[0], kernel_size=7)
+        self.conv2 = downsample_conv(conv_planes[0], conv_planes[1], kernel_size=5)
+        self.conv3 = downsample_conv(conv_planes[1], conv_planes[2])
+        self.conv4 = downsample_conv(conv_planes[2], conv_planes[3])
+        self.conv5 = downsample_conv(conv_planes[3], conv_planes[4])
+        self.conv6 = downsample_conv(conv_planes[4], conv_planes[5])
+        self.conv7 = downsample_conv(conv_planes[5], conv_planes[6])
+
+        upconv_planes = [512, 512, 256, 128, 64, 32, 16]
+        self.upconv7 = upconv(conv_planes[6],   upconv_planes[0])
+        self.upconv6 = upconv(upconv_planes[0], upconv_planes[1])
+        self.upconv5 = upconv(upconv_planes[1], upconv_planes[2])
+        self.upconv4 = upconv(upconv_planes[2], upconv_planes[3])
+        self.upconv3 = upconv(upconv_planes[3], upconv_planes[4])
+        self.upconv2 = upconv(upconv_planes[4], upconv_planes[5])
+        self.upconv1 = upconv(upconv_planes[5], upconv_planes[6])
+
+        self.iconv7 = conv(upconv_planes[0] + conv_planes[5], upconv_planes[0])
+        self.iconv6 = conv(upconv_planes[1] + conv_planes[4], upconv_planes[1])
+        self.iconv5 = conv(upconv_planes[2] + conv_planes[3], upconv_planes[2])
+        self.iconv4 = conv(upconv_planes[3] + conv_planes[2], upconv_planes[3])
+        self.iconv3 = conv(1 + upconv_planes[4] + conv_planes[1], upconv_planes[4])
+        self.iconv2 = conv(1 + upconv_planes[5] + conv_planes[0], upconv_planes[5])
+        self.iconv1 = conv(1 + upconv_planes[6], upconv_planes[6])
+
+        self.predict_disp4 = predict_disp(upconv_planes[3])
+        self.predict_disp3 = predict_disp(upconv_planes[4])
+        self.predict_disp2 = predict_disp(upconv_planes[5])
+        self.predict_disp1 = predict_disp(upconv_planes[6])
+
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+                xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    zeros_(m.bias)
+
+    def forward(self, x):
+        out_conv1 = self.conv1(x)
+        out_conv2 = self.conv2(out_conv1)
+        out_conv3 = self.conv3(out_conv2)
+        out_conv4 = self.conv4(out_conv3)
+        out_conv5 = self.conv5(out_conv4)
+        out_conv6 = self.conv6(out_conv5)
+        out_conv7 = self.conv7(out_conv6)
+
+        out_upconv7 = crop_like(self.upconv7(out_conv7), out_conv6)
+        concat7 = torch.cat((out_upconv7, out_conv6), 1)
+        out_iconv7 = self.iconv7(concat7)
+
+        out_upconv6 = crop_like(self.upconv6(out_iconv7), out_conv5)
+        concat6 = torch.cat((out_upconv6, out_conv5), 1)
+        out_iconv6 = self.iconv6(concat6)
+
+        out_upconv5 = crop_like(self.upconv5(out_iconv6), out_conv4)
+        concat5 = torch.cat((out_upconv5, out_conv4), 1)
+        out_iconv5 = self.iconv5(concat5)
+
+        out_upconv4 = crop_like(self.upconv4(out_iconv5), out_conv3)
+        concat4 = torch.cat((out_upconv4, out_conv3), 1)
+        out_iconv4 = self.iconv4(concat4)
+        disp4 = self.alpha * self.predict_disp4(out_iconv4) + self.beta
+
+        out_upconv3 = crop_like(self.upconv3(out_iconv4), out_conv2)
+        disp4_up = crop_like(F.interpolate(disp4, scale_factor=2, mode='bilinear', align_corners=False), out_conv2)
+        concat3 = torch.cat((out_upconv3, out_conv2, disp4_up), 1)
+        out_iconv3 = self.iconv3(concat3)
+        disp3 = self.alpha * self.predict_disp3(out_iconv3) + self.beta
+
+        out_upconv2 = crop_like(self.upconv2(out_iconv3), out_conv1)
+        disp3_up = crop_like(F.interpolate(disp3, scale_factor=2, mode='bilinear', align_corners=False), out_conv1)
+        concat2 = torch.cat((out_upconv2, out_conv1, disp3_up), 1)
+        out_iconv2 = self.iconv2(concat2)
+        disp2 = self.alpha * self.predict_disp2(out_iconv2) + self.beta
+
+        out_upconv1 = crop_like(self.upconv1(out_iconv2), x)
+        disp2_up = crop_like(F.interpolate(disp2, scale_factor=2, mode='bilinear', align_corners=False), x)
+        concat1 = torch.cat((out_upconv1, disp2_up), 1)
+        out_iconv1 = self.iconv1(concat1)
+        disp1 = self.alpha * self.predict_disp1(out_iconv1) + self.beta
+
+        return disp1, disp2, disp3, disp4
+        

spikezoo/archs/stir/package_core/package_core/flow_utils.py

@@ -0,0 +1,167 @@
+import cv2
+import sys
+import numpy as np
+import argparse
+
+def load_flow(path):
+    with open(path, 'rb') as f:
+        magic = float(np.fromfile(f, np.float32, count = 1)[0])
+        if magic == 202021.25:
+            w, h = np.fromfile(f, np.int32, count = 1)[0], np.fromfile(f, np.int32, count = 1)[0]
+            data = np.fromfile(f, np.float32, count = h*w*2)
+            data.resize((h, w, 2))
+            return data
+        return None
+
+def save_flow(path, flow):
+    magic = np.array([202021.25], np.float32)
+    h, w = flow.shape[:2]
+    h, w = np.array([h], np.int32), np.array([w], np.int32)
+
+    with open(path, 'wb') as f:
+        magic.tofile(f); w.tofile(f); h.tofile(f); flow.tofile(f)
+
+def makeColorwheel():
+
+	#  color encoding scheme
+
+	#   adapted from the color circle idea described at
+	#   http://members.shaw.ca/quadibloc/other/colint.htm
+
+	RY = 15
+	YG = 6
+	GC = 4
+	CB = 11
+	BM = 13
+	MR = 6
+
+	ncols = RY + YG + GC + CB + BM + MR
+
+	colorwheel = np.zeros([ncols, 3]) # r g b
+
+	col = 0
+	#RY
+	colorwheel[0:RY, 0] = 255
+	colorwheel[0:RY, 1] = np.floor(255*np.arange(0, RY, 1)/RY)
+	col += RY
+
+	#YG
+	colorwheel[col:YG+col, 0]= 255 - np.floor(255*np.arange(0, YG, 1)/YG)
+	colorwheel[col:YG+col, 1] = 255;
+	col += YG;
+
+	#GC
+	colorwheel[col:GC+col, 1]= 255 
+	colorwheel[col:GC+col, 2] = np.floor(255*np.arange(0, GC, 1)/GC)
+	col += GC;
+
+	#CB
+	colorwheel[col:CB+col, 1]= 255 - np.floor(255*np.arange(0, CB, 1)/CB)
+	colorwheel[col:CB+col, 2] = 255
+	col += CB;
+
+	#BM
+	colorwheel[col:BM+col, 2]= 255 
+	colorwheel[col:BM+col, 0] = np.floor(255*np.arange(0, BM, 1)/BM)
+	col += BM;
+
+	#MR
+	colorwheel[col:MR+col, 2]= 255 - np.floor(255*np.arange(0, MR, 1)/MR)
+	colorwheel[col:MR+col, 0] = 255
+	return 	colorwheel
+
+def computeColor(u, v):
+
+	colorwheel = makeColorwheel();
+	nan_u = np.isnan(u)
+	nan_v = np.isnan(v)
+	nan_u = np.where(nan_u)
+	nan_v = np.where(nan_v) 
+
+	u[nan_u] = 0
+	u[nan_v] = 0
+	v[nan_u] = 0 
+	v[nan_v] = 0
+
+	ncols = colorwheel.shape[0]
+	radius = np.sqrt(u**2 + v**2)
+	a = np.arctan2(-v, -u) / np.pi
+	fk = (a+1) /2 * (ncols-1) # -1~1 maped to 1~ncols
+	k0 = fk.astype(np.uint8)	 # 1, 2, ..., ncols
+	k1 = k0+1
+	k1[k1 == ncols] = 0
+	f = fk - k0
+
+	img = np.empty([k1.shape[0], k1.shape[1],3])
+	ncolors = colorwheel.shape[1]
+	for i in range(ncolors):
+		tmp = colorwheel[:,i]
+		col0 = tmp[k0]/255
+		col1 = tmp[k1]/255
+		col = (1-f)*col0 + f*col1
+		idx = radius <= 1
+		col[idx] = 1 - radius[idx]*(1-col[idx]) # increase saturation with radius    
+		col[~idx] *= 0.75 # out of range
+		img[:,:,2-i] = np.floor(255*col).astype(np.uint8)
+
+	return img.astype(np.uint8)
+
+
+def flow2rgb(flow):
+	# H, W, 2
+	eps = sys.float_info.epsilon
+	UNKNOWN_FLOW_THRESH = 1e9
+	UNKNOWN_FLOW = 1e10
+
+	u = flow[:,:,0]
+	v = flow[:,:,1]
+
+	maxu = -999
+	maxv = -999
+
+	minu = 999
+	minv = 999
+
+	maxrad = -1
+	#fix unknown flow
+	greater_u = np.where(u > UNKNOWN_FLOW_THRESH)
+	greater_v = np.where(v > UNKNOWN_FLOW_THRESH)
+	u[greater_u] = 0
+	u[greater_v] = 0
+	v[greater_u] = 0 
+	v[greater_v] = 0
+
+	maxu = max([maxu, np.amax(u)])
+	minu = min([minu, np.amin(u)])
+
+	maxv = max([maxv, np.amax(v)])
+	minv = min([minv, np.amin(v)])
+	rad = np.sqrt(np.multiply(u,u)+np.multiply(v,v))
+	maxrad = max([maxrad, np.amax(rad)])
+	# print('max flow: %.4f flow range: u = %.3f .. %.3f; v = %.3f .. %.3f\n' % (maxrad, minu, maxu, minv, maxv))
+
+	u = u/(maxrad+eps)
+	v = v/(maxrad+eps)
+	img = computeColor(u, v)
+	return img[:,:,[2,1,0]]
+
+def flow_to_numpy_rgb(flow):
+    flow_map_np = flow.detach().cpu().numpy().transpose((0, 2, 3, 1))
+    B, H, W, _ = flow_map_np.shape
+    colored_rgb = np.empty([B, H, W, 3])
+    for i in range(B):
+        colored_array= flow2rgb(flow_map_np[i])
+        colored_rgb[i] = colored_array 
+    return colored_rgb.astype(np.uint8)
+
+#if __name__ == '__main__':
+#    import matplotlib.pyplot as plt#
+#
+#    flow = load_flow('/home/autovision/mycode/PWC-Net/PyTorch/tmp/frame_0010.flo')
+#    #flow = load_flow('datasets/Sintel/training/flow/alley_1/frame_0001.flo')
+#    img = vis_flow(flow)
+#    import imageio
+#    imageio.imsave('test.png', img)
+#    import cv2
+#    cv2.imshow('', img[:,:,:])
+#    cv2.waitKey()

spikezoo/archs/stir/package_core/package_core/generic_train_test.py

@@ -0,0 +1,76 @@
+import os
+import time
+from skimage import io
+
+from .metrics import *
+from .image_proc import *
+
+class Generic_train_test():
+	def __init__(self, model, opts, dataloader, logger, dataloader_val=None):
+		self.model=model
+		self.opts=opts
+		self.dataloader=dataloader
+		self.logger=logger
+		self.dataloader_val = dataloader_val
+
+	def decode_input(self, data):
+		raise NotImplementedError()
+
+	def validation(self):
+		raise NotImplementedError()
+
+	def train_single_iterate(self, data, total_steps, epoch):
+		_input=self.decode_input(data)
+
+		self.model.set_input(_input)
+		self.model.optimize_parameters()
+
+		#=========== visualize results ============#
+		if total_steps % self.opts.log_freq==0:
+			info = self.model.get_current_scalars()
+			for tag, value in info.items():
+				self.logger.add_scalar(tag, value, total_steps)
+
+			results = self.model.get_current_visuals()
+			for tag, images in results.items():
+				self.logger.add_images(tag, images, total_steps)
+
+			print('epoch', epoch, 'steps', total_steps)
+			print('losses', info)
+
+	def train(self):
+		total_steps = 0
+		if self.dataloader is not None:
+			print('#training images ', len(self.dataloader)*self.opts.batch_sz)
+
+		for epoch in range(self.opts.start_epoch, self.opts.max_epochs):
+			if epoch > self.opts.lr_start_epoch_decay - self.opts.lr_step:
+				self.model.update_lr()
+
+			if epoch % self.opts.save_freq==0 or epoch <= self.opts.save_begin:
+				self.model.save_checkpoint(str(epoch))
+
+			if self.dataloader is not None:
+				for i, data in enumerate(self.dataloader):
+					total_steps+=1
+					self.train_single_iterate(data, total_steps, epoch)
+			else:
+				for i in range(10000):
+					total_steps+=1
+					self.train_single_iterate(None, total_steps, epoch)
+
+			# validation if dataloader provided
+			if self.dataloader_val is not None:
+				self.validation(epoch)
+
+	def train_single_instance(self):
+		total_steps = 0
+		data=iter(self.dataloader).next()
+
+		for epoch in range(10000):
+			for i in range(1000):
+				total_steps+=1
+				self.train_single_iterate(data, total_steps, epoch)
+				
+
+