#coding:utf-8
import cv2
import numpy as np
import imutils
import matplotlib.pyplot as plt
import os
from multiprocessing import Pool,Process
import time
def get_dis(A,B):
BT = np.transpose(B)
A_BT = np.dot(A, BT)
Asq = A ** 2
Asq = np.tile(np.sum(Asq, axis=1, keepdims=True), (1, A_BT.shape[1]))
Bsq = BT ** 2
Bsq = np.tile(np.sum(Bsq, axis=0, keepdims=True), (A_BT.shape[0], 1))
ED = np.sqrt(Asq + Bsq - 2 * A_BT)
# print('ED.shape',ED.shape)
long_axis_ind = np.unravel_index(np.argmax(ED), ED.shape)
return ED[long_axis_ind[0],long_axis_ind[1]],long_axis_ind
def show_points(axis_ind,points1,points2,color):
i_ = axis_ind[0]
j_ = axis_ind[1]
x1, y1 = points1[i_, :]
x2, y2 = points2[j_, :]
plt.plot([x1, x2], [y1, y2], color=color)
return x1, y1,x2, y2
def get_contour_centroid(cnt):
m = cv2.moments(cnt)
c_x = int(m['m10'] / m['m00'])
c_y = int(m['m01'] / m['m00'])
return c_x, c_y
def get_line_sides(points,center_x):
side_points=[]
another_side_points = []
for point in points:
if point[0]>center_x:
side_points.append(point)
else:
another_side_points.append(point)
return np.array(side_points),np.array(another_side_points)
def get_disk_axis(image_list_path):
# image_list_path='./data/sichuan_pig_mistake_label/2018-9-11-IMG_0319.jpg'
print(image_list_path)
img=cv2.imread(image_list_path)
img=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
h,w=img.shape
image_thre = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)[1]
cnts = cv2.findContours(image_thre, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if imutils.is_cv2() else cnts[1]
c = sorted(cnts, key=cv2.contourArea, reverse=True)
disk_cnt = np.squeeze(c[0])
# plt.plot(disk_cnt[:,0],disk_cnt[:,1])
mask=np.zeros(shape=(h,w,3))
dummy_mask = cv2.drawContours(mask, [disk_cnt], 0, (255, 0, 0))
y, x = np.where(dummy_mask[:, :, 0] == 255)
plt.figure()
plt.gca().set_aspect('equal')
disk_complet_points = np.stack((x, y), axis=-1)
center_x,center_y=get_contour_centroid(disk_cnt)
# print('center_x,center_y=',center_x,center_y)
plt.plot(disk_complet_points[:, 0], disk_complet_points[:, 1], 'o')
###
plt.scatter(center_x,center_y,color='black')
side_points,another_side_points=get_line_sides(disk_complet_points,center_x=center_x)
# print('disk_complet_points.shape',disk_complet_points.shape)
# print('side_points.shape',side_points.shape)
# print('another_side_points.shape', another_side_points.shape)
long_axis_length,long_axis_ind=get_dis(side_points,another_side_points)
print('long_axis_length=',long_axis_length)
x1, y1, x2, y2=show_points(long_axis_ind,side_points,another_side_points,color='red')
del_side_points = np.delete(side_points, obj=(long_axis_ind[0]), axis=0)
del_another_side_points = np.delete(another_side_points, obj=(long_axis_ind[1]), axis=0)
del_complet_points = np.vstack((del_side_points, del_another_side_points))
short_axis_points = []
if y2 - y1:
k= -(x2-x1)/(y2-y1)
b = center_y - k * center_x
print('k,b',k,b)
for point in del_complet_points:
if abs(np.ceil(point[1] - k * point[0] - b)) < 20:
short_axis_points.append(point)
else:
for point in del_complet_points:
if abs(center_x-point[0])<10:
short_axis_points.append(point)
short_side_points, short_another_side_points = get_line_sides(np.array(short_axis_points), center_x=center_x)
short_axis_length, short_axis_ind = get_dis(short_side_points, short_another_side_points)
_,_,_,_,=show_points(short_axis_ind,short_side_points, short_another_side_points,color='blue')
# plt.show()
#debug see output
out_path='./data/sichuan_pig_mistake_label/disk_mask_out'
if not os.path.exists(out_path):
os.mkdir(out_path)
plt.savefig(out_path+'/'+image_list_path.split('/')[-1])
def main():
start_time = time.time()
pool = Pool()
path='./data/sichuan_pig_mistake_label/disk_mask'
images_list_path=[os.path.join(path,i) for i in os.listdir(path)]
for count,image_list_path in enumerate(images_list_path):
pool.apply_async(func=get_disk_axis,args=(image_list_path,))
pool.close()
pool.join()
print("time={}".format(time.time() - start_time))
if __name__ == '__main__':
main()
红点代表长轴点,黑点代表短轴点。