Fitting best possible two separate lines to a set of 2D points with python-CodePudding

I have a black-white image where I want to fit 2 separate lines to the edges in the image. It is easy to fit a single line with opencv, with the code below. How do I fit two best possible lines to this image. Here is the input image and 1 line result. I need something that is like the last image below.

def fit_line_to_edges(edges):
    points = np.where(edges>0)
    points = np.array([points[1], points[0]]).T
    [vx,vy,x,y] = cv2.fitLine(points, cv2.DIST_L2,0,0.01,0.01)
    lefty = int((-x*vy/vx)   y)
    righty = int(((edges.shape[1]-x)*vy/vx) y)
    cv2.line(edges, (edges.shape[1]-1,righty), (0,lefty), color = (136,155,112), thickness = 3)
    return edges

CodePudding user response：

Thanks for the help in the comments. Here is the working code using a piecewise linear function

from scipy import optimize
import matplotlib.pyplot as plt
import numpy as np
import cv2

def fit_piecewise_lines(edges):
    def piecewise_linear(x, x0, y0, k1, k2):
        return np.piecewise(x, [x < x0], [lambda x:k1*x   y0-k1*x0, lambda x:k2*x   y0-k2*x0])

    # get x and y points from edges
    points = np.where(edges>0)
    points = np.array([points[1], points[0]]).T
    x = points[:,0]
    y = points[:,1]

    # initial point
    p_init = [np.mean(x), np.mean(y), 1, 1]

    # set initial parameter estimates
    p, e = optimize.curve_fit(piecewise_linear, x, y, p0=p_init)

    # get the point locations
    pt1 = (np.min(x).astype(int), piecewise_linear(np.min(x), *p).astype(int))
    pt3 = (np.max(x).astype(int), piecewise_linear(np.max(x), *p).astype(int))
    pt2 = (p[0].astype(int), p[1].astype(int))

    # plot it
    cv2.line(edges, pt1, pt2, (255,0,0), 5)
    cv2.line(edges, pt2, pt3, (0,255,0), 5)


image_file = "./edges.png"
edges = cv2.imread(image_file)

fit_piecewise_lines(edges)

plt.figure(figsize=(10,10))
plt.imshow(edges)