#!/usr/bin/env python # coding: utf-8 import numpy as np import os import sys import copy import time import argparse from skimage.io import imread, imsave from skimage.color import rgb2gray from skimage import exposure from skimage.transform import resize from PIL import Image, ImageDraw, ImageColor def gradient4D(img): (row, col) = img.shape g4d = np.zeros((row, col)) for i in range(row-1): for j in range(col-1): g4d[i, j] = abs(img[i+1, j] - img[i, j] ) + abs(img[i, j+1] - img[i, j]) + abs(img[i+1, j-1] - img[i, j]) + abs(img[i+1, j+1] - img[i, j]) return npNormalise(g4d) def npNormalise(xArray): XNorm = (xArray - xArray.min()) / (xArray.max() - xArray.min()) return XNorm def getBaseLines(ws): baseLineList = [] baseLineListIdex = [] for i in range(0,ws-5): #if i == 1: break for j in range(5,ws): # cut bord #if j == 6: break # 1 # ------------- # | | # 4 | |2 # | | # | | # -------------- # 3 # adjacent edge img12 = Image.new('F', (ws,ws),0) draw12 = ImageDraw.Draw(img12) draw12.line(xy=(i, 0, ws-1, j), fill=(1), width = 1) baseLineList.append(np.asarray(img12)) baseLineListIdex.append(np.asarray([[i, 0],[ws-1, j]])) baseLineList.append(np.rot90(np.asarray(img12), 1, axes=(0, 1))) baseLineListIdex.append(np.asarray([[j, 0],[0, ws-1-i]])) baseLineList.append(np.rot90(np.asarray(img12), 2, axes=(0, 1))) baseLineListIdex.append(np.asarray([[0, ws-1-j],[ws-1-i, ws-1]])) baseLineList.append(np.rot90(np.asarray(img12), 3, axes=(0, 1))) baseLineListIdex.append(np.asarray([[ws-1, i],[ws-1-j, ws-1]])) # opposite side img13 = Image.new('F', (ws,ws),0) draw13 = ImageDraw.Draw(img13) draw13.line(xy=(i, 0, j, ws-1), fill=(1), width = 1) baseLineList.append(np.asarray(img13)) baseLineListIdex.append(np.asarray([[i,0],[j, ws-1]])) baseLineList.append(np.asarray(img13).T) baseLineListIdex.append(np.asarray([[0,i],[ws-1, j]])) print('base line number :', len(baseLineList)) return np.asarray(baseLineList), np.asarray(baseLineListIdex) def calculSlope(v1,v2): difX = v2[0] - v1[0] difY = v2[1] - v1[1] if difX == 0 : lk = 5*difY else: lk = difY / difX return lk def clusterRegion(centerLine, scale = 4, windowSize=64): H = windowSize W = windowSize sMask = np.zeros([H,W]) ix = int(centerLine[0][0]) iy = int(centerLine[0][1]) pixelRange = int(min(H,W) / scale) # scale = 10 k = calculSlope(centerLine[0],centerLine[1]) if abs(k) > 1: while ix > 0: iy = int(round(((ix-centerLine[0][1]) / k) + centerLine[0][0])) frontY = max(0, iy-pixelRange) backY = min(W,iy+pixelRange+1) sMask[ix, frontY:backY] = 1 ix = ix - 1 ix = int(centerLine[0][0]) while ix < H: iy = int(round(((ix-centerLine[0][1]) / k) + centerLine[0][0])) frontY = max(0, iy-pixelRange) backY = min(W,iy+pixelRange+1) sMask[ix, frontY:backY] = 1 ix = ix + 1 else: while iy > 0: ix = int(round(((iy-centerLine[0][0]) * k) + centerLine[0][1])) frontX = max(0, ix-pixelRange) backX = min(H,ix+pixelRange+1) sMask[frontX:backX, iy] = 1 iy = iy - 1 iy = int(centerLine[0][1]) while iy < W: ix = int(round(((iy-centerLine[0][0]) * k) + centerLine[0][1])) frontX = max(0, ix-pixelRange) backX = min(H,ix+pixelRange+1) sMask[frontX:backX, iy] = 1 iy = iy + 1 return sMask def drawGroupLine(file, lineList, flineListCluster, scale, functionName, colorSTR, outputPath): c = ImageColor.colormap cList = list(c.items()) (inputPath,inputFile) = os.path.split(file) print(inputPath) print(inputFile) with Image.open(file) as img4draw: w, h = img4draw.size scale = 1/64 wScale = np.ceil(w*scale) hScale = np.ceil(h*scale) img1 = ImageDraw.Draw(img4draw) # for n,lineSet in enumerate(flineListCluster): # for [v1,v2],w,_ in lineSet[1:]: # img1.line([(v1[0]*wScale,v1[1]*hScale), (v2[0]*wScale,v2[1]*hScale)], fill = cList[int(n*2)+2][1], width = 4) for [v1,v2] in lineList: img1.line([(v1[0],v1[1]), (v2[0],v2[1])], fill = colorSTR, width = 8) img4draw.save(os.path.join(outputPath, inputFile[:-4] + '_' + str(functionName) + inputFile[-4:] )) def sortSlope(lineListArray): slopeList = [] groupWeight = 0 for l in lineListArray: if (l[0][1][0] - l[0][0][0] ) == 0: k = 1000 else: k = (l[0][1][1] - l[0][0][1]) / (l[0][1][0] - l[0][0][0]) slopeList.append(k) groupWeight = groupWeight + l[1] # print('weight = ', l[1]) index = np.argsort(np.array(slopeList)) #groupWeight = np.mean(groupWeight) return [lineListArray[int(np.median(index))][0], lineListArray[0][1], lineListArray[int(np.median(index))][2]] # index[len(index)//2] def forceLinesClusterIntegration(cluster): forceL = [] for i,lineSet in enumerate(cluster): forceL.append(sortSlope(lineSet[1:])) return forceL def refine(lineList, fg, wg, iP, ws): wlist = [] forceList = [] for l in lineList: wlist.append(l[1]) npwList = np.array(wlist) sortWeight = npwList.argsort()[::-1] for n,wId in enumerate(sortWeight): if n == 0: gMask = clusterRegion(lineList[wId][0], fg, ws) forceList.append([gMask, lineList[wId]]) else: judge, forceList = judgeVertexAdvanced(lineList[wId][2], lineList[wId][0], npwList[wId], forceList, wg, iP ) if judge == False: gMask = clusterRegion(lineList[wId][0], fg, ws) forceList.append([gMask, lineList[wId]]) flList = forceLinesClusterIntegration(forceList) return flList def findSaliantLineCluster(gradient4d,allLines,allLinesIndex,ws, orgW, orgH ): weightList = [] fineGrained0 = 8 intePrec0 = 0.8 forceLinesCluster = [] for l in allLines: w = np.sum(gradient4d*l) weightList.append(w) npWeightList = np.array(weightList) sortWeightList = npWeightList.argsort()[::-1] # [::-1] inverse a list for n,wId in enumerate(sortWeightList[:300]): if n == 0: groupMask = clusterRegion(allLinesIndex[wId], fineGrained0, ws) forceLinesCluster.append([groupMask, [allLinesIndex[wId], npWeightList[wId], allLines[wId]]]) # else: # print(npWeightList[sortWeightList[n-1]]) # print(npWeightList[wId]) if (npWeightList[sortWeightList[n-1]] - npWeightList[wId]) > 10 : print('weight break------in line ', str(n)) break judge, forceLinesCluster = judgeVertexAdvanced(allLines[wId], allLinesIndex[wId], npWeightList[wId], forceLinesCluster , 2, intePrec0) if judge == False: groupMask = clusterRegion(allLinesIndex[wId], fineGrained0, ws) forceLinesCluster.append([groupMask, [allLinesIndex[wId], npWeightList[wId], allLines[wId]]]) forceLinesRough = forceLinesClusterIntegration(forceLinesCluster) forceLinesRoughNew = forceLinesRough forceLinesRoughOrg = [] fineGrained = 7 wGrained = 3 intePrec = 0.7 for i in range(10000): if len(forceLinesRoughNew) == len(forceLinesRoughOrg): if (fineGrained <= 4 )and (wGrained >= 10) : print('break in loop ', str(i)) break forceLinesRoughOrg = forceLinesRoughNew forceLinesRoughNew = refine(forceLinesRoughNew, fineGrained, wGrained, intePrec, ws) if fineGrained > 4: fineGrained = fineGrained-1 if intePrec > 0.6: intePrec = intePrec - 0.05 if wGrained < 10: wGrained = wGrained + 1 forceLines = [] for l in forceLinesRoughNew: forceLines.append(l[0]) forceLines = np.array(forceLines) scale = 1/ws HWscale = np.array([[np.ceil(orgW*scale),np.ceil(orgH*scale)], [np.ceil(orgW*scale),np.ceil(orgH*scale)]]) HWS = np.expand_dims(HWscale,0).repeat(forceLines.shape[0],axis=0) forceLines = forceLines*HWS return forceLines, forceLinesCluster,HWS def judgeVertexAdvanced(line1,v1, v1w, forceL, wSeuil = 4, intersectPrecent = 0.7): v1 = np.array(v1) newGroup = False for cl in forceL: vPossible = cl[0]*line1 if np.sum(vPossible) > (np.sum(line1)*intersectPrecent): if abs(cl[1][1] - v1w) < wSeuil: cl.append([v1,v1w,line1]) return True,forceL else: return True,forceL return False, forceL def getLeadingLine(imgpath, outPath): windowSize = 64 allLines, allLinesIndex = getBaseLines(windowSize) img = imread(imgpath) print(img.shape) if (len(img.shape) != 3) or (img.shape[2] != 3): print('NOT a 3 channel image') else: orgH, orgW, _ = img.shape resizeImg = resize(img,(windowSize,windowSize)) # Add contrast logImg = exposure.adjust_log(resizeImg, 1) grayImg = rgb2gray(logImg) gradient4d= gradient4D(grayImg) forceLines, forceLinesCluster, scale = findSaliantLineCluster(gradient4d,allLines,allLinesIndex,windowSize, orgW, orgH ) drawGroupLine(imgpath, forceLines, forceLinesCluster, scale, 'forceLines', 'red', outPath) if __name__ == '__main__': parser = argparse.ArgumentParser(description='Find the Probable leading lines, please provide 1) your input image path or a folder path for input images, and 2) the output folder you wish.') parser.add_argument('input', type=str, help='The path for your input image or folder') parser.add_argument('-o', '--output', type=str, default='./OUTPUT', help='The path for your output folder ') args = parser.parse_args() INPUT_DIRECTORY = args.input OUTPUT_DIRECTORY = args.output print('INPUT : ', INPUT_DIRECTORY) print('OUTPUT : ', OUTPUT_DIRECTORY) if not (os.path.exists(OUTPUT_DIRECTORY)): print('Create output path:' , OUTPUT_DIRECTORY) os.makedirs(OUTPUT_DIRECTORY) start = time.time() if os.path.isfile( INPUT_DIRECTORY ): if INPUT_DIRECTORY.lower().endswith(('.jpg', '.png')) and not INPUT_DIRECTORY.lower().startswith('.'): getLeadingLine(INPUT_DIRECTORY,OUTPUT_DIRECTORY) elif os.path.isdir( INPUT_DIRECTORY ): files= os.listdir(INPUT_DIRECTORY) for i, file in enumerate(files): if file.lower().endswith(('.jpg', '.png')) and not file.lower().startswith('.'): fullpath = os.path.join(INPUT_DIRECTORY, file) getLeadingLine(fullpath,OUTPUT_DIRECTORY) end = time.time() print(' use time = ', str((end - start)/60.0), 'm')