''' This script is helper function for preprocessing. Most of the code are converted from LayoutNet official's matlab code. All functions, naming rule and data flow follow official for easier converting and comparing. Code is not optimized for python or numpy yet. ''' import sys import numpy as np from scipy.ndimage import map_coordinates import cv2 from pylsd import lsd def computeUVN(n, in_, planeID): ''' compute v given u and normal. ''' if planeID == 2: n = np.array([n[1], n[2], n[0]]) elif planeID == 3: n = np.array([n[2], n[0], n[1]]) bc = n[0] * np.sin(in_) + n[1] * np.cos(in_) bs = n[2] out = np.arctan(-bc / (bs + 1e-9)) return out def computeUVN_vec(n, in_, planeID): ''' vectorization version of computeUVN @n N x 3 @in_ MN x 1 @planeID N ''' n = n.copy() if (planeID == 2).sum(): n[planeID == 2] = np.roll(n[planeID == 2], 2, axis=1) if (planeID == 3).sum(): n[planeID == 3] = np.roll(n[planeID == 3], 1, axis=1) n = np.repeat(n, in_.shape[0] // n.shape[0], axis=0) assert n.shape[0] == in_.shape[0] bc = n[:, [0]] * np.sin(in_) + n[:, [1]] * np.cos(in_) bs = n[:, [2]] out = np.arctan(-bc / (bs + 1e-9)) return out def xyz2uvN(xyz, planeID=1): ID1 = (int(planeID) - 1 + 0) % 3 ID2 = (int(planeID) - 1 + 1) % 3 ID3 = (int(planeID) - 1 + 2) % 3 normXY = np.sqrt(xyz[:, [ID1]] ** 2 + xyz[:, [ID2]] ** 2) normXY[normXY < 0.000001] = 0.000001 normXYZ = np.sqrt(xyz[:, [ID1]] ** 2 + xyz[:, [ID2]] ** 2 + xyz[:, [ID3]] ** 2) v = np.arcsin(xyz[:, [ID3]] / normXYZ) u = np.arcsin(xyz[:, [ID1]] / normXY) valid = (xyz[:, [ID2]] < 0) & (u >= 0) u[valid] = np.pi - u[valid] valid = (xyz[:, [ID2]] < 0) & (u <= 0) u[valid] = -np.pi - u[valid] uv = np.hstack([u, v]) uv[np.isnan(uv[:, 0]), 0] = 0 return uv def uv2xyzN(uv, planeID=1): ID1 = (int(planeID) - 1 + 0) % 3 ID2 = (int(planeID) - 1 + 1) % 3 ID3 = (int(planeID) - 1 + 2) % 3 xyz = np.zeros((uv.shape[0], 3)) xyz[:, ID1] = np.cos(uv[:, 1]) * np.sin(uv[:, 0]) xyz[:, ID2] = np.cos(uv[:, 1]) * np.cos(uv[:, 0]) xyz[:, ID3] = np.sin(uv[:, 1]) return xyz def uv2xyzN_vec(uv, planeID): ''' vectorization version of uv2xyzN @uv N x 2 @planeID N ''' assert (planeID.astype(int) != planeID).sum() == 0 planeID = planeID.astype(int) ID1 = (planeID - 1 + 0) % 3 ID2 = (planeID - 1 + 1) % 3 ID3 = (planeID - 1 + 2) % 3 ID = np.arange(len(uv)) xyz = np.zeros((len(uv), 3)) xyz[ID, ID1] = np.cos(uv[:, 1]) * np.sin(uv[:, 0]) xyz[ID, ID2] = np.cos(uv[:, 1]) * np.cos(uv[:, 0]) xyz[ID, ID3] = np.sin(uv[:, 1]) return xyz def warpImageFast(im, XXdense, YYdense): minX = max(1., np.floor(XXdense.min()) - 1) minY = max(1., np.floor(YYdense.min()) - 1) maxX = min(im.shape[1], np.ceil(XXdense.max()) + 1) maxY = min(im.shape[0], np.ceil(YYdense.max()) + 1) im = im[int(round(minY-1)):int(round(maxY)), int(round(minX-1)):int(round(maxX))] assert XXdense.shape == YYdense.shape out_shape = XXdense.shape coordinates = [ (YYdense - minY).reshape(-1), (XXdense - minX).reshape(-1), ] im_warp = np.stack([ map_coordinates(im[..., c], coordinates, order=1).reshape(out_shape) for c in range(im.shape[-1])], axis=-1) return im_warp def rotatePanorama(img, vp=None, R=None): ''' Rotate panorama if R is given, vp (vanishing point) will be overlooked otherwise R is computed from vp ''' sphereH, sphereW, C = img.shape # new uv coordinates TX, TY = np.meshgrid(range(1, sphereW + 1), range(1, sphereH + 1)) TX = TX.reshape(-1, 1, order='F') TY = TY.reshape(-1, 1, order='F') ANGx = (TX - sphereW/2 - 0.5) / sphereW * np.pi * 2 ANGy = -(TY - sphereH/2 - 0.5) / sphereH * np.pi uvNew = np.hstack([ANGx, ANGy]) xyzNew = uv2xyzN(uvNew, 1) # rotation matrix if R is None: R = np.linalg.inv(vp.T) xyzOld = np.linalg.solve(R, xyzNew.T).T uvOld = xyz2uvN(xyzOld, 1) Px = (uvOld[:, 0] + np.pi) / (2*np.pi) * sphereW + 0.5 Py = (-uvOld[:, 1] + np.pi/2) / np.pi * sphereH + 0.5 Px = Px.reshape(sphereH, sphereW, order='F') Py = Py.reshape(sphereH, sphereW, order='F') # boundary imgNew = np.zeros((sphereH+2, sphereW+2, C), np.float64) imgNew[1:-1, 1:-1, :] = img imgNew[1:-1, 0, :] = img[:, -1, :] imgNew[1:-1, -1, :] = img[:, 0, :] imgNew[0, 1:sphereW//2+1, :] = img[0, sphereW-1:sphereW//2-1:-1, :] imgNew[0, sphereW//2+1:-1, :] = img[0, sphereW//2-1::-1, :] imgNew[-1, 1:sphereW//2+1, :] = img[-1, sphereW-1:sphereW//2-1:-1, :] imgNew[-1, sphereW//2+1:-1, :] = img[0, sphereW//2-1::-1, :] imgNew[0, 0, :] = img[0, 0, :] imgNew[-1, -1, :] = img[-1, -1, :] imgNew[0, -1, :] = img[0, -1, :] imgNew[-1, 0, :] = img[-1, 0, :] rotImg = warpImageFast(imgNew, Px+1, Py+1) return rotImg def imgLookAt(im, CENTERx, CENTERy, new_imgH, fov): sphereH = im.shape[0] sphereW = im.shape[1] warped_im = np.zeros((new_imgH, new_imgH, 3)) TX, TY = np.meshgrid(range(1, new_imgH + 1), range(1, new_imgH + 1)) TX = TX.reshape(-1, 1, order='F') TY = TY.reshape(-1, 1, order='F') TX = TX - 0.5 - new_imgH/2 TY = TY - 0.5 - new_imgH/2 r = new_imgH / 2 / np.tan(fov/2) # convert to 3D R = np.sqrt(TY ** 2 + r ** 2) ANGy = np.arctan(- TY / r) ANGy = ANGy + CENTERy X = np.sin(ANGy) * R Y = -np.cos(ANGy) * R Z = TX INDn = np.nonzero(np.abs(ANGy) > np.pi/2) # project back to sphere ANGx = np.arctan(Z / -Y) RZY = np.sqrt(Z ** 2 + Y ** 2) ANGy = np.arctan(X / RZY) ANGx[INDn] = ANGx[INDn] + np.pi ANGx = ANGx + CENTERx INDy = np.nonzero(ANGy < -np.pi/2) ANGy[INDy] = -np.pi - ANGy[INDy] ANGx[INDy] = ANGx[INDy] + np.pi INDx = np.nonzero(ANGx <= -np.pi); ANGx[INDx] = ANGx[INDx] + 2 * np.pi INDx = np.nonzero(ANGx > np.pi); ANGx[INDx] = ANGx[INDx] - 2 * np.pi INDx = np.nonzero(ANGx > np.pi); ANGx[INDx] = ANGx[INDx] - 2 * np.pi INDx = np.nonzero(ANGx > np.pi); ANGx[INDx] = ANGx[INDx] - 2 * np.pi Px = (ANGx + np.pi) / (2*np.pi) * sphereW + 0.5 Py = ((-ANGy) + np.pi/2) / np.pi * sphereH + 0.5 INDxx = np.nonzero(Px < 1) Px[INDxx] = Px[INDxx] + sphereW im = np.concatenate([im, im[:, :2]], 1) Px = Px.reshape(new_imgH, new_imgH, order='F') Py = Py.reshape(new_imgH, new_imgH, order='F') warped_im = warpImageFast(im, Px, Py) return warped_im def separatePano(panoImg, fov, x, y, imgSize=320): '''cut a panorama image into several separate views''' assert x.shape == y.shape if not isinstance(fov, np.ndarray): fov = fov * np.ones_like(x) sepScene = [ { 'img': imgLookAt(panoImg.copy(), xi, yi, imgSize, fovi), 'vx': xi, 'vy': yi, 'fov': fovi, 'sz': imgSize, } for xi, yi, fovi in zip(x, y, fov) ] return sepScene def lsdWrap(img): ''' Opencv implementation of Rafael Grompone von Gioi, Jérémie Jakubowicz, Jean-Michel Morel, and Gregory Randall, LSD: a Line Segment Detector, Image Processing On Line, vol. 2012. [Rafael12] http://www.ipol.im/pub/art/2012/gjmr-lsd/?utm_source=doi @img input image ''' if len(img.shape) == 3: img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) lines = lsd(img, quant=0.7) if lines is None: return np.zeros_like(img), np.array([]) edgeMap = np.zeros_like(img) for i in range(lines.shape[0]): pt1 = (int(lines[i, 0]), int(lines[i, 1])) pt2 = (int(lines[i, 2]), int(lines[i, 3])) width = lines[i, 4] cv2.line(edgeMap, pt1, pt2, 255, int(np.ceil(width / 2))) edgeList = np.concatenate([lines, np.ones_like(lines[:, :2])], 1) return edgeMap, edgeList def edgeFromImg2Pano(edge): edgeList = edge['edgeLst'] if len(edgeList) == 0: return np.array([]) vx = edge['vx'] vy = edge['vy'] fov = edge['fov'] imH, imW = edge['img'].shape R = (imW/2) / np.tan(fov/2) # im is the tangent plane, contacting with ball at [x0 y0 z0] x0 = R * np.cos(vy) * np.sin(vx) y0 = R * np.cos(vy) * np.cos(vx) z0 = R * np.sin(vy) vecposX = np.array([np.cos(vx), -np.sin(vx), 0]) vecposY = np.cross(np.array([x0, y0, z0]), vecposX) vecposY = vecposY / np.sqrt(vecposY @ vecposY.T) vecposX = vecposX.reshape(1, -1) vecposY = vecposY.reshape(1, -1) Xc = (0 + imW-1) / 2 Yc = (0 + imH-1) / 2 vecx1 = edgeList[:, [0]] - Xc vecy1 = edgeList[:, [1]] - Yc vecx2 = edgeList[:, [2]] - Xc vecy2 = edgeList[:, [3]] - Yc vec1 = np.tile(vecx1, [1, 3]) * vecposX + np.tile(vecy1, [1, 3]) * vecposY vec2 = np.tile(vecx2, [1, 3]) * vecposX + np.tile(vecy2, [1, 3]) * vecposY coord1 = [[x0, y0, z0]] + vec1 coord2 = [[x0, y0, z0]] + vec2 normal = np.cross(coord1, coord2, axis=1) normal = normal / np.linalg.norm(normal, axis=1, keepdims=True) panoList = np.hstack([normal, coord1, coord2, edgeList[:, [-1]]]) return panoList def _intersection(range1, range2): if range1[1] < range1[0]: range11 = [range1[0], 1] range12 = [0, range1[1]] else: range11 = range1 range12 = [0, 0] if range2[1] < range2[0]: range21 = [range2[0], 1] range22 = [0, range2[1]] else: range21 = range2 range22 = [0, 0] b = max(range11[0], range21[0]) < min(range11[1], range21[1]) if b: return b b2 = max(range12[0], range22[0]) < min(range12[1], range22[1]) b = b or b2 return b def _insideRange(pt, range): if range[1] > range[0]: b = pt >= range[0] and pt <= range[1] else: b1 = pt >= range[0] and pt <= 1 b2 = pt >= 0 and pt <= range[1] b = b1 or b2 return b def combineEdgesN(edges): ''' Combine some small line segments, should be very conservative OUTPUT lines: combined line segments ori_lines: original line segments line format [nx ny nz projectPlaneID umin umax LSfov score] ''' arcList = [] for edge in edges: panoLst = edge['panoLst'] if len(panoLst) == 0: continue arcList.append(panoLst) arcList = np.vstack(arcList) # ori lines numLine = len(arcList) ori_lines = np.zeros((numLine, 8)) areaXY = np.abs(arcList[:, 2]) areaYZ = np.abs(arcList[:, 0]) areaZX = np.abs(arcList[:, 1]) planeIDs = np.argmax(np.stack([areaXY, areaYZ, areaZX], -1), 1) + 1 # XY YZ ZX for i in range(numLine): ori_lines[i, :3] = arcList[i, :3] ori_lines[i, 3] = planeIDs[i] coord1 = arcList[i, 3:6] coord2 = arcList[i, 6:9] uv = xyz2uvN(np.stack([coord1, coord2]), planeIDs[i]) umax = uv[:, 0].max() + np.pi umin = uv[:, 0].min() + np.pi if umax - umin > np.pi: ori_lines[i, 4:6] = np.array([umax, umin]) / 2 / np.pi else: ori_lines[i, 4:6] = np.array([umin, umax]) / 2 / np.pi ori_lines[i, 6] = np.arccos(( np.dot(coord1, coord2) / (np.linalg.norm(coord1) * np.linalg.norm(coord2)) ).clip(-1, 1)) ori_lines[i, 7] = arcList[i, 9] # additive combination lines = ori_lines.copy() for _ in range(3): numLine = len(lines) valid_line = np.ones(numLine, bool) for i in range(numLine): if not valid_line[i]: continue dotProd = (lines[:, :3] * lines[[i], :3]).sum(1) valid_curr = np.logical_and((np.abs(dotProd) > np.cos(np.pi / 180)), valid_line) valid_curr[i] = False for j in np.nonzero(valid_curr)[0]: range1 = lines[i, 4:6] range2 = lines[j, 4:6] valid_rag = _intersection(range1, range2) if not valid_rag: continue # combine I = np.argmax(np.abs(lines[i, :3])) if lines[i, I] * lines[j, I] > 0: nc = lines[i, :3] * lines[i, 6] + lines[j, :3] * lines[j, 6] else: nc = lines[i, :3] * lines[i, 6] - lines[j, :3] * lines[j, 6] nc = nc / np.linalg.norm(nc) if _insideRange(range1[0], range2): nrmin = range2[0] else: nrmin = range1[0] if _insideRange(range1[1], range2): nrmax = range2[1] else: nrmax = range1[1] u = np.array([[nrmin], [nrmax]]) * 2 * np.pi - np.pi v = computeUVN(nc, u, lines[i, 3]) xyz = uv2xyzN(np.hstack([u, v]), lines[i, 3]) l = np.arccos(np.dot(xyz[0, :], xyz[1, :]).clip(-1, 1)) scr = (lines[i,6]*lines[i,7] + lines[j,6]*lines[j,7]) / (lines[i,6]+lines[j,6]) lines[i] = [*nc, lines[i, 3], nrmin, nrmax, l, scr] valid_line[j] = False lines = lines[valid_line] return lines, ori_lines def icosahedron2sphere(level): # this function use a icosahedron to sample uniformly on a sphere a = 2 / (1 + np.sqrt(5)) M = np.array([ 0, a, -1, a, 1, 0, -a, 1, 0, 0, a, 1, -a, 1, 0, a, 1, 0, 0, a, 1, 0, -a, 1, -1, 0, a, 0, a, 1, 1, 0, a, 0, -a, 1, 0, a, -1, 0, -a, -1, 1, 0, -a, 0, a, -1, -1, 0, -a, 0, -a, -1, 0, -a, 1, a, -1, 0, -a, -1, 0, 0, -a, -1, -a, -1, 0, a, -1, 0, -a, 1, 0, -1, 0, a, -1, 0, -a, -a, -1, 0, -1, 0, -a, -1, 0, a, a, 1, 0, 1, 0, -a, 1, 0, a, a, -1, 0, 1, 0, a, 1, 0, -a, 0, a, 1, -1, 0, a, -a, 1, 0, 0, a, 1, a, 1, 0, 1, 0, a, 0, a, -1, -a, 1, 0, -1, 0, -a, 0, a, -1, 1, 0, -a, a, 1, 0, 0, -a, -1, -1, 0, -a, -a, -1, 0, 0, -a, -1, a, -1, 0, 1, 0, -a, 0, -a, 1, -a, -1, 0, -1, 0, a, 0, -a, 1, 1, 0, a, a, -1, 0]) coor = M.T.reshape(3, 60, order='F').T coor, idx = np.unique(coor, return_inverse=True, axis=0) tri = idx.reshape(3, 20, order='F').T # extrude coor = list(coor / np.tile(np.linalg.norm(coor, axis=1, keepdims=True), (1, 3))) for _ in range(level): triN = [] for t in range(len(tri)): n = len(coor) coor.append((coor[tri[t, 0]] + coor[tri[t, 1]]) / 2) coor.append((coor[tri[t, 1]] + coor[tri[t, 2]]) / 2) coor.append((coor[tri[t, 2]] + coor[tri[t, 0]]) / 2) triN.append([n, tri[t, 0], n+2]) triN.append([n, tri[t, 1], n+1]) triN.append([n+1, tri[t, 2], n+2]) triN.append([n, n+1, n+2]) tri = np.array(triN) # uniquefy coor, idx = np.unique(coor, return_inverse=True, axis=0) tri = idx[tri] # extrude coor = list(coor / np.tile(np.sqrt(np.sum(coor * coor, 1, keepdims=True)), (1, 3))) return np.array(coor), np.array(tri) def curveFitting(inputXYZ, weight): ''' @inputXYZ: N x 3 @weight : N x 1 ''' l = np.linalg.norm(inputXYZ, axis=1, keepdims=True) inputXYZ = inputXYZ / l weightXYZ = inputXYZ * weight XX = np.sum(weightXYZ[:, 0] ** 2) YY = np.sum(weightXYZ[:, 1] ** 2) ZZ = np.sum(weightXYZ[:, 2] ** 2) XY = np.sum(weightXYZ[:, 0] * weightXYZ[:, 1]) YZ = np.sum(weightXYZ[:, 1] * weightXYZ[:, 2]) ZX = np.sum(weightXYZ[:, 2] * weightXYZ[:, 0]) A = np.array([ [XX, XY, ZX], [XY, YY, YZ], [ZX, YZ, ZZ]]) U, S, Vh = np.linalg.svd(A) outputNM = Vh[-1, :] outputNM = outputNM / np.linalg.norm(outputNM) return outputNM def sphereHoughVote(segNormal, segLength, segScores, binRadius, orthTolerance, candiSet, force_unempty=True): # initial guess numLinesg = len(segNormal) voteBinPoints = candiSet.copy() voteBinPoints = voteBinPoints[~(voteBinPoints[:,2] < 0)] reversValid = (segNormal[:, 2] < 0).reshape(-1) segNormal[reversValid] = -segNormal[reversValid] voteBinUV = xyz2uvN(voteBinPoints) numVoteBin = len(voteBinPoints) voteBinValues = np.zeros(numVoteBin) for i in range(numLinesg): tempNorm = segNormal[[i]] tempDots = (voteBinPoints * tempNorm).sum(1) valid = np.abs(tempDots) < np.cos((90 - binRadius) * np.pi / 180) voteBinValues[valid] = voteBinValues[valid] + segScores[i] * segLength[i] checkIDs1 = np.nonzero(voteBinUV[:, [1]] > np.pi / 3)[0] voteMax = 0 checkID1Max = 0 checkID2Max = 0 checkID3Max = 0 for j in range(len(checkIDs1)): checkID1 = checkIDs1[j] vote1 = voteBinValues[checkID1] if voteBinValues[checkID1] == 0 and force_unempty: continue checkNormal = voteBinPoints[[checkID1]] dotProduct = (voteBinPoints * checkNormal).sum(1) checkIDs2 = np.nonzero(np.abs(dotProduct) < np.cos((90 - orthTolerance) * np.pi / 180))[0] for i in range(len(checkIDs2)): checkID2 = checkIDs2[i] if voteBinValues[checkID2] == 0 and force_unempty: continue vote2 = vote1 + voteBinValues[checkID2] cpv = np.cross(voteBinPoints[checkID1], voteBinPoints[checkID2]).reshape(1, 3) cpn = np.linalg.norm(cpv) dotProduct = (voteBinPoints * cpv).sum(1) / cpn checkIDs3 = np.nonzero(np.abs(dotProduct) > np.cos(orthTolerance * np.pi / 180))[0] for k in range(len(checkIDs3)): checkID3 = checkIDs3[k] if voteBinValues[checkID3] == 0 and force_unempty: continue vote3 = vote2 + voteBinValues[checkID3] if vote3 > voteMax: lastStepCost = vote3 - voteMax if voteMax != 0: tmp = (voteBinPoints[[checkID1Max, checkID2Max, checkID3Max]] * \ voteBinPoints[[checkID1, checkID2, checkID3]]).sum(1) lastStepAngle = np.arccos(tmp.clip(-1, 1)) else: lastStepAngle = np.zeros(3) checkID1Max = checkID1 checkID2Max = checkID2 checkID3Max = checkID3 voteMax = vote3 if checkID1Max == 0: print('[WARN] sphereHoughVote: no orthogonal voting exist', file=sys.stderr) return None, 0, 0 initXYZ = voteBinPoints[[checkID1Max, checkID2Max, checkID3Max]] # refine refiXYZ = np.zeros((3, 3)) dotprod = (segNormal * initXYZ[[0]]).sum(1) valid = np.abs(dotprod) < np.cos((90 - binRadius) * np.pi / 180) validNm = segNormal[valid] validWt = segLength[valid] * segScores[valid] validWt = validWt / validWt.max() refiNM = curveFitting(validNm, validWt) refiXYZ[0] = refiNM.copy() dotprod = (segNormal * initXYZ[[1]]).sum(1) valid = np.abs(dotprod) < np.cos((90 - binRadius) * np.pi / 180) validNm = segNormal[valid] validWt = segLength[valid] * segScores[valid] validWt = validWt / validWt.max() validNm = np.vstack([validNm, refiXYZ[[0]]]) validWt = np.vstack([validWt, validWt.sum(0, keepdims=1) * 0.1]) refiNM = curveFitting(validNm, validWt) refiXYZ[1] = refiNM.copy() refiNM = np.cross(refiXYZ[0], refiXYZ[1]) refiXYZ[2] = refiNM / np.linalg.norm(refiNM) return refiXYZ, lastStepCost, lastStepAngle def findMainDirectionEMA(lines): '''compute vp from set of lines''' # initial guess segNormal = lines[:, :3] segLength = lines[:, [6]] segScores = np.ones((len(lines), 1)) shortSegValid = (segLength < 5 * np.pi / 180).reshape(-1) segNormal = segNormal[~shortSegValid, :] segLength = segLength[~shortSegValid] segScores = segScores[~shortSegValid] numLinesg = len(segNormal) candiSet, tri = icosahedron2sphere(3) ang = np.arccos((candiSet[tri[0,0]] * candiSet[tri[0,1]]).sum().clip(-1, 1)) / np.pi * 180 binRadius = ang / 2 initXYZ, score, angle = sphereHoughVote(segNormal, segLength, segScores, 2*binRadius, 2, candiSet) if initXYZ is None: print('[WARN] findMainDirectionEMA: initial failed', file=sys.stderr) return None, score, angle # iterative refine iter_max = 3 candiSet, tri = icosahedron2sphere(5) numCandi = len(candiSet) angD = np.arccos((candiSet[tri[0, 0]] * candiSet[tri[0, 1]]).sum().clip(-1, 1)) / np.pi * 180 binRadiusD = angD / 2 curXYZ = initXYZ.copy() tol = np.linspace(4*binRadius, 4*binRadiusD, iter_max) # shrink down ls and candi for it in range(iter_max): dot1 = np.abs((segNormal * curXYZ[[0]]).sum(1)) dot2 = np.abs((segNormal * curXYZ[[1]]).sum(1)) dot3 = np.abs((segNormal * curXYZ[[2]]).sum(1)) valid1 = dot1 < np.cos((90 - tol[it]) * np.pi / 180) valid2 = dot2 < np.cos((90 - tol[it]) * np.pi / 180) valid3 = dot3 < np.cos((90 - tol[it]) * np.pi / 180) valid = valid1 | valid2 | valid3 if np.sum(valid) == 0: print('[WARN] findMainDirectionEMA: zero line segments for voting', file=sys.stderr) break subSegNormal = segNormal[valid] subSegLength = segLength[valid] subSegScores = segScores[valid] dot1 = np.abs((candiSet * curXYZ[[0]]).sum(1)) dot2 = np.abs((candiSet * curXYZ[[1]]).sum(1)) dot3 = np.abs((candiSet * curXYZ[[2]]).sum(1)) valid1 = dot1 > np.cos(tol[it] * np.pi / 180) valid2 = dot2 > np.cos(tol[it] * np.pi / 180) valid3 = dot3 > np.cos(tol[it] * np.pi / 180) valid = valid1 | valid2 | valid3 if np.sum(valid) == 0: print('[WARN] findMainDirectionEMA: zero line segments for voting', file=sys.stderr) break subCandiSet = candiSet[valid] tcurXYZ, _, _ = sphereHoughVote(subSegNormal, subSegLength, subSegScores, 2*binRadiusD, 2, subCandiSet) if tcurXYZ is None: print('[WARN] findMainDirectionEMA: no answer found', file=sys.stderr) break curXYZ = tcurXYZ.copy() mainDirect = curXYZ.copy() mainDirect[0] = mainDirect[0] * np.sign(mainDirect[0,2]) mainDirect[1] = mainDirect[1] * np.sign(mainDirect[1,2]) mainDirect[2] = mainDirect[2] * np.sign(mainDirect[2,2]) uv = xyz2uvN(mainDirect) I1 = np.argmax(uv[:,1]) J = np.setdiff1d(np.arange(3), I1) I2 = np.argmin(np.abs(np.sin(uv[J,0]))) I2 = J[I2] I3 = np.setdiff1d(np.arange(3), np.hstack([I1, I2])) mainDirect = np.vstack([mainDirect[I1], mainDirect[I2], mainDirect[I3]]) mainDirect[0] = mainDirect[0] * np.sign(mainDirect[0,2]) mainDirect[1] = mainDirect[1] * np.sign(mainDirect[1,1]) mainDirect[2] = mainDirect[2] * np.sign(mainDirect[2,0]) mainDirect = np.vstack([mainDirect, -mainDirect]) return mainDirect, score, angle def multi_linspace(start, stop, num): div = (num - 1) y = np.arange(0, num, dtype=np.float64) steps = (stop - start) / div return steps.reshape(-1, 1) * y + start.reshape(-1, 1) def assignVanishingType(lines, vp, tol, area=10): numLine = len(lines) numVP = len(vp) typeCost = np.zeros((numLine, numVP)) # perpendicular for vid in range(numVP): cosint = (lines[:, :3] * vp[[vid]]).sum(1) typeCost[:, vid] = np.arcsin(np.abs(cosint).clip(-1, 1)) # infinity u = np.stack([lines[:, 4], lines[:, 5]], -1) u = u.reshape(-1, 1) * 2 * np.pi - np.pi v = computeUVN_vec(lines[:, :3], u, lines[:, 3]) xyz = uv2xyzN_vec(np.hstack([u, v]), np.repeat(lines[:, 3], 2)) xyz = multi_linspace(xyz[0::2].reshape(-1), xyz[1::2].reshape(-1), 100) xyz = np.vstack([blk.T for blk in np.split(xyz, numLine)]) xyz = xyz / np.linalg.norm(xyz, axis=1, keepdims=True) for vid in range(numVP): ang = np.arccos(np.abs((xyz * vp[[vid]]).sum(1)).clip(-1, 1)) notok = (ang < area * np.pi / 180).reshape(numLine, 100).sum(1) != 0 typeCost[notok, vid] = 100 I = typeCost.min(1) tp = typeCost.argmin(1) tp[I > tol] = numVP + 1 return tp, typeCost def refitLineSegmentB(lines, vp, vpweight=0.1): ''' Refit direction of line segments INPUT: lines: original line segments vp: vannishing point vpweight: if set to 0, lines will not change; if set to inf, lines will be forced to pass vp ''' numSample = 100 numLine = len(lines) xyz = np.zeros((numSample+1, 3)) wei = np.ones((numSample+1, 1)) wei[numSample] = vpweight * numSample lines_ali = lines.copy() for i in range(numLine): n = lines[i, :3] sid = lines[i, 4] * 2 * np.pi eid = lines[i, 5] * 2 * np.pi if eid < sid: x = np.linspace(sid, eid + 2 * np.pi, numSample) % (2 * np.pi) else: x = np.linspace(sid, eid, numSample) u = -np.pi + x.reshape(-1, 1) v = computeUVN(n, u, lines[i, 3]) xyz[:numSample] = uv2xyzN(np.hstack([u, v]), lines[i, 3]) xyz[numSample] = vp outputNM = curveFitting(xyz, wei) lines_ali[i, :3] = outputNM return lines_ali def paintParameterLine(parameterLine, width, height): lines = parameterLine.copy() panoEdgeC = np.zeros((height, width)) num_sample = max(height, width) for i in range(len(lines)): n = lines[i, :3] sid = lines[i, 4] * 2 * np.pi eid = lines[i, 5] * 2 * np.pi if eid < sid: x = np.linspace(sid, eid + 2 * np.pi, num_sample) x = x % (2 * np.pi) else: x = np.linspace(sid, eid, num_sample) u = -np.pi + x.reshape(-1, 1) v = computeUVN(n, u, lines[i, 3]) xyz = uv2xyzN(np.hstack([u, v]), lines[i, 3]) uv = xyz2uvN(xyz, 1) m = np.minimum(np.floor((uv[:,0] + np.pi) / (2 * np.pi) * width) + 1, width).astype(np.int32) n = np.minimum(np.floor(((np.pi / 2) - uv[:, 1]) / np.pi * height) + 1, height).astype(np.int32) panoEdgeC[n-1, m-1] = i return panoEdgeC def panoEdgeDetection(img, viewSize=320, qError=0.7, refineIter=3): ''' line detection on panorama INPUT: img: image waiting for detection, double type, range 0~1 viewSize: image size of croped views qError: set smaller if more line segment wanted OUTPUT: oLines: detected line segments vp: vanishing point views: separate views of panorama edges: original detection of line segments in separate views panoEdge: image for visualize line segments ''' cutSize = viewSize fov = np.pi / 3 xh = np.arange(-np.pi, np.pi*5/6, np.pi/6) yh = np.zeros(xh.shape[0]) xp = np.array([-3/3, -2/3, -1/3, 0/3, 1/3, 2/3, -3/3, -2/3, -1/3, 0/3, 1/3, 2/3]) * np.pi yp = np.array([ 1/4, 1/4, 1/4, 1/4, 1/4, 1/4, -1/4, -1/4, -1/4, -1/4, -1/4, -1/4]) * np.pi x = np.concatenate([xh, xp, [0, 0]]) y = np.concatenate([yh, yp, [np.pi/2., -np.pi/2]]) sepScene = separatePano(img.copy(), fov, x, y, cutSize) edge = [] for i, scene in enumerate(sepScene): edgeMap, edgeList = lsdWrap(scene['img']) edge.append({ 'img': edgeMap, 'edgeLst': edgeList, 'vx': scene['vx'], 'vy': scene['vy'], 'fov': scene['fov'], }) edge[-1]['panoLst'] = edgeFromImg2Pano(edge[-1]) lines, olines = combineEdgesN(edge) clines = lines.copy() for _ in range(refineIter): mainDirect, score, angle = findMainDirectionEMA(clines) tp, typeCost = assignVanishingType(lines, mainDirect[:3], 0.1, 10) lines1 = lines[tp==0] lines2 = lines[tp==1] lines3 = lines[tp==2] lines1rB = refitLineSegmentB(lines1, mainDirect[0], 0) lines2rB = refitLineSegmentB(lines2, mainDirect[1], 0) lines3rB = refitLineSegmentB(lines3, mainDirect[2], 0) clines = np.vstack([lines1rB, lines2rB, lines3rB]) panoEdge1r = paintParameterLine(lines1rB, img.shape[1], img.shape[0]) panoEdge2r = paintParameterLine(lines2rB, img.shape[1], img.shape[0]) panoEdge3r = paintParameterLine(lines3rB, img.shape[1], img.shape[0]) panoEdger = np.stack([panoEdge1r, panoEdge2r, panoEdge3r], -1) # output olines = clines vp = mainDirect views = sepScene edges = edge panoEdge = panoEdger return olines, vp, views, edges, panoEdge, score, angle if __name__ == '__main__': # disable OpenCV3's non thread safe OpenCL option cv2.ocl.setUseOpenCL(False) import os import argparse import PIL from PIL import Image import time parser = argparse.ArgumentParser() parser.add_argument('--i', required=True) parser.add_argument('--o_prefix', required=True) parser.add_argument('--qError', default=0.7, type=float) parser.add_argument('--refineIter', default=3, type=int) args = parser.parse_args() # Read image img_ori = np.array(Image.open(args.i).resize((1024, 512))) # Vanishing point estimation & Line segments detection s_time = time.time() olines, vp, views, edges, panoEdge, score, angle = panoEdgeDetection(img_ori, qError=args.qError, refineIter=args.refineIter) print('Elapsed time: %.2f' % (time.time() - s_time)) panoEdge = (panoEdge > 0) print('Vanishing point:') for v in vp[2::-1]: print('%.6f %.6f %.6f' % tuple(v)) # Visualization edg = rotatePanorama(panoEdge.astype(np.float64), vp[2::-1]) img = rotatePanorama(img_ori / 255.0, vp[2::-1]) one = img.copy() * 0.5 one[(edg > 0.5).sum(-1) > 0] = 0 one[edg[..., 0] > 0.5, 0] = 1 one[edg[..., 1] > 0.5, 1] = 1 one[edg[..., 2] > 0.5, 2] = 1 Image.fromarray((edg * 255).astype(np.uint8)).save('%s_edg.png' % args.o_prefix) Image.fromarray((img * 255).astype(np.uint8)).save('%s_img.png' % args.o_prefix) Image.fromarray((one * 255).astype(np.uint8)).save('%s_one.png' % args.o_prefix)