leading_lines/LigneForce.py

#Copyright (C) [2023] [ZHANG Jing, Université du Littoral Côte d'Opale]
#
#This program is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with this program.  If not, see <http://www.gnu.org/licenses/>.

#!/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

## compute image gradient map, take the absolute difference values in 4 dimensions for each pixel
  
#        --------------------------
#       |                          |
#       |          i,  j   i,  j+1 |
#       | i+1,j-1  i+1,j   i+1,j+1 |
#        --------------------------
  
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)

# normalise values to [0, 1]
def npNormalise(xArray):
    XNorm = (xArray - xArray.min()) / (xArray.max() - xArray.min())
    return XNorm

# compute all potential lines for a square of size ws*ws
# input : length of the square
# output : 1. a set of binary images, each image contains only one line
#           2. a set containing the coordinates of the start points and the end points of the line
def getBaseLines(ws):
    baseLineList = []
    baseLineListIdex = []
    # Skip 5 pixels to avoid lines near edges
    for i in range(0,ws-5):
        for j in range(5,ws):  #  cut bord
            
    #               1
    #         -------------
    #         |           |
    #       4 |           |2
    #         |           |
    #         |           |
    #         --------------
    #               3
            # adjacent edge   (like edge 1 and edge 2, edge 1 and edge 4)
            img12 = Image.new('F', (ws,ws),0)
            draw12 = ImageDraw.Draw(img12)
            # lines that the start point in edge 1 and the end point in edge 2
            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]]))
            # lines that the start point in edge 4 and the end point in edge 1
            baseLineList.append(np.rot90(np.asarray(img12), 1, axes=(0, 1)))
            baseLineListIdex.append(np.asarray([[j, 0],[0, ws-1-i]]))
            # lines that the start point in edge 3 and the end point in edge 4
            baseLineList.append(np.rot90(np.asarray(img12), 2, axes=(0, 1)))
            baseLineListIdex.append(np.asarray([[0, ws-1-j],[ws-1-i, ws-1]]))
            # lines that the start point in edge 2 and the end point in edge 3
            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)
            # lines that the start point in edge 4 and the end point in edge 2
            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]]))
            # lines that the start point in edge 1 and the end point in edge 3
            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)

# Calculate the slope of the line formed by vertex1 and vertex2
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

# Compute the band mask of a line
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])
    # calculate the width of band mask
    pixelRange = int(min(H,W) / scale)  #  scale = 10 
    
    # get the slope of line
    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


# fonction for display all the lines
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)
    
    # read the orignal file for draw
    with Image.open(file) as img4draw:
        w, h = img4draw.size
        if w >h:     # add lineWidth to adapt the visibility of drawing results to different image sizes
            lineWidth = int(h/40)
        else:
            lineWidth = int(w/40)
            
        scale = 1/64
        wScale = np.ceil(w*scale)
        hScale = np.ceil(h*scale)
        
        img1 = ImageDraw.Draw(img4draw) 
        
        # draw the cluster result
#         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)
        # draw all the centers
        for [v1,v2] in lineList:
            img1.line([(v1[0],v1[1]), (v2[0],v2[1])], fill = colorSTR, width = lineWidth)
            
        img4draw.save(os.path.join(outputPath, inputFile[:-4] + '_' + str(functionName) + inputFile[-4:] ))

# sort the slope of lines, inutile for version 0
def sortSlope(lineListArray):
    print('lineListArray', 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])
    print('slopeList : ', slopeList)
    index = np.argsort(np.array(slopeList))
    print('sortSlope index : ', index)
    print('sortSlope index median : ', int(np.median(index)))
    #groupWeight = np.mean(groupWeight)
    
    return [lineListArray[int(np.median(index))][0], lineListArray[0][1], lineListArray[int(np.median(index))][2]]
    # return [lineListArray[int(len(index)/2)][0], lineListArray[0][1], lineListArray[int(len(index)/2)][2]]
#     index[len(index)//2]

# extraction the center of each group
def forceLinesClusterIntegration(cluster):
    forceL = []
    
    for i,lineSet in enumerate(cluster):
        forceL.append(lineSet[1])
    return forceL


# refine the cluster result
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     
                
# the main process of clustering the lines
def findSaliantLineCluster(gradient4d,allLines,allLinesIndex,ws, orgW, orgH ):
    weightList = []
    fineGrained0 = 8   # initial refine grained = ws / 8
    intePrec0 = 0.8    # initial intersection precision
    forceLinesCluster = []
    
    # compute weights of lines
    for l in allLines:
        w = np.sum(gradient4d*l)
        weightList.append(w)

    npWeightList = np.array(weightList)
    sortWeightList = npWeightList.argsort()[::-1]   # [::-1] inverse a list, range from large to small
    # top 300 weighted candidates, about 0.14% of the total lines
    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')