kitti深度图补全

rgb图像

这个是雷达稀疏图

 这个是补全后的图像

code后续整理出来放（python）

import scipy
import skimage
import os
import numpy as np
# from pypardiso import spsolve
from scipy.sparse.linalg import spsolve
from PIL import Image
from skimage.color import rgb2gray
import matplotlib.pyplot as plt
import matplotlib
from matplotlib import cm
import multiprocessing as mp
from tqdm import tqdm
import time
import numpngw


def fill_depth_colorization(imgRgb=None, imgDepthInput=None, alpha=1):
	imgIsNoise = imgDepthInput == 0
	maxImgAbsDepth = np.max(imgDepthInput)
	imgDepth = imgDepthInput / maxImgAbsDepth
	imgDepth[imgDepth > 1] = 1
	(H, W) = imgDepth.shape
	numPix = H * W
	indsM = np.arange(numPix).reshape((W, H)).transpose()
	knownValMask = (imgIsNoise == False).astype(int)
	grayImg = rgb2gray(imgRgb)
	winRad = 1
	len_ = 0
	absImgNdx = 0
	len_window = (2 * winRad + 1) ** 2
	len_zeros = numPix * len_window

	cols = np.zeros(len_zeros) - 1
	rows = np.zeros(len_zeros) - 1
	vals = np.zeros(len_zeros) - 1
	gvals = np.zeros(len_window) - 1

	for j in range(W):
		for i in range(H):
			nWin = 0
			for ii in range(max(0, i - winRad), min(i + winRad + 1, H)):
				for jj in range(max(0, j - winRad), min(j + winRad + 1, W)):
					if ii == i and jj == j:
						continue

					rows[len_] = absImgNdx
					cols[len_] = indsM[ii, jj]
					gvals[nWin] = grayImg[ii, jj]

					len_ = len_ + 1
					nWin = nWin + 1

			curVal = grayImg[i, j]
			gvals[nWin] = curVal
			c_var = np.mean((gvals[:nWin + 1] - np.mean(gvals[:nWin+ 1])) ** 2)

			csig = c_var * 0.6
			mgv = np.min((gvals[:nWin] - curVal) ** 2)
			if csig < -mgv / np.log(0.01):
				csig = -mgv / np.log(0.01)

			if csig < 2e-06:
				csig = 2e-06

			gvals[:nWin] = np.exp(-(gvals[:nWin] - curVal) ** 2 / csig)
			gvals[:nWin] = gvals[:nWin] / sum(gvals[:nWin])
			vals[len_ - nWin:len_] = -gvals[:nWin]
			rows[len_] = absImgNdx
			cols[len_] = absImgNdx
			vals[len_] = 1  # sum(gvals(1:nWin))
			len_ = len_ + 1
			absImgNdx = absImgNdx + 1
	vals = vals[:len_]
	cols = cols[:len_]
	rows = rows[:len_]
	A = scipy.sparse.csr_matrix((vals, (rows, cols)), (numPix, numPix))
	rows = np.arange(0, numPix)
	cols = np.arange(0, numPix)
	vals = (knownValMask * alpha).transpose().reshape(numPix)
	G = scipy.sparse.csr_matrix((vals, (rows, cols)), (numPix, numPix))
	A = A + G
	b = np.multiply(vals.reshape(numPix), imgDepth.flatten('F'))
	new_vals = spsolve(A, b)
	new_vals = np.reshape(new_vals, (H, W), 'F')
	denoisedDepthImg = new_vals * maxImgAbsDepth
	output = denoisedDepthImg.reshape((H, W)).astype('float32')
	output = np.multiply(output, (1-knownValMask)) + imgDepthInput
	return output


def main():
	
	dir_rgb = 'C:\\Users\\xaut_\\Desktop\\complete\\rgb'			# depth_2011_09_26_drive_0001_sync_image_02_0000000005
	dir_depth_imple = 'C:\\Users\\xaut_\\Desktop\\complete\\depth_imple'
	after_comple = 'C:\\Users\\xaut_\\Desktop\\complete\\after_comple'			# depth_2011_09_26_drive_0001_sync_image_02_0000000005.png
	num = 0
	for rgb in tqdm(sorted(os.listdir(dir_rgb))):
		start = time.time()
		if rgb not in os.listdir(after_comple):
			rgb_path = os.path.join(dir_rgb, rgb)
			depth_path = os.path.join(dir_depth_imple, rgb)
			img = Image.open(rgb_path)
			depth = Image.open(depth_path)
			try:
				img = np.array(img, dtype=np.float32)
				depth = np.array(depth, dtype=np.float32)
				out = fill_depth_colorization(img/255., depth)
			except:
				print('*********本张图片有问题**********')
				continue
			out = Image.fromarray(np.uint16(out))
			numpngw.write_png('C:\\Users\\xaut_\\Desktop\\complete\\after_comple\\{}'.format(rgb), out)   # 保存成16位


			num +=1
			end = time.time()
			print('已完成{}张，还剩{}张, 本次花费{:.2f}秒'.format(len(os.listdir(after_comple)), str(len(os.listdir(dir_rgb))-len(os.listdir(after_comple))), (end-start)))
			
			
if __name__ == '__main__':
	print('开始')
	main()

平均补全一张用时23s左右，你有好的优化方法可以一起交流（kitti需要补全大概3w左右图片）

觉得有用点个赞

上色之后的补全图：