Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import cv2 | |
| def draw_output(out, img): | |
| img = img.copy() | |
| # print(out['lf'][0].shape[0], out['sol'].shape[0]) | |
| # sys.exit() | |
| for i in range(out['sol'].shape[0]): | |
| j=i | |
| p = out['sol'][i] | |
| c = int(255 * p[-1]) | |
| color = (c,0,255-c) | |
| x = p[0] | |
| y = p[1] | |
| r = p[2] | |
| x_comp = np.cos(r) | |
| y_comp = -np.sin(r) | |
| s = p[3] | |
| rx = x + s * x_comp * 2 | |
| ry = y + s * y_comp * 2 | |
| rx2 = x - s * x_comp | |
| ry2 = y - s * y_comp | |
| rx = int(rx) | |
| ry = int(ry) | |
| rx2 = int(rx2) | |
| ry2 = int(ry2) | |
| x = int(x) | |
| y = int(y) | |
| scale = abs(int(s)) | |
| color = (0,0,255) | |
| cv2.circle(img,(x,y), int(scale), color, 2) | |
| cv2.circle(img,(x,y), 4, color, -1) | |
| cv2.arrowedLine(img, (x,y), (rx,ry), color, 2, tipLength=0.25) | |
| # cv2.line(img, (rx2,ry2), (rx,ry), color, 2) | |
| cv2.putText(img,str(i),(x,y), cv2.FONT_HERSHEY_SIMPLEX, 1,(0,255,0),2,cv2.LINE_AA) | |
| # for j in range(out['lf'][0].shape[0]): | |
| begin = out['beginning'][j] | |
| end = out['ending'][j] | |
| last_xy = None | |
| # for i in xrange(len(out['lf'])): | |
| begin_f = int(np.floor(begin)) | |
| end_f = int(np.ceil(end)) | |
| for i in range(begin_f, end_f+1): | |
| if i == begin_f: | |
| p0 = out['lf'][i][j].mean(axis=1) | |
| p1 = out['lf'][i+1][j].mean(axis=1) | |
| t = begin - np.floor(begin) | |
| p = p0 * (1 - t) + p1 * t | |
| elif i == end_f: | |
| p0 = out['lf'][i-1][j].mean(axis=1) | |
| if i != len(out['lf']): | |
| p1 = out['lf'][i][j].mean(axis=1) | |
| t = end - np.floor(end) | |
| p = p0 * (1 - t) + p1 * t | |
| else: | |
| p = p0 | |
| else: | |
| p = out['lf'][i][j].mean(axis=1) | |
| x = p[0] | |
| y = p[1] | |
| x = int(x) | |
| y = int(y) | |
| # c = int(255 * p[-1]) | |
| # color = (c,0,255-c) | |
| color = (0,150,0) | |
| cv2.circle(img,(x,y), 4, color, -1) | |
| if last_xy is not None: | |
| cv2.line(img, (x,y), last_xy, color, int(s)) | |
| last_xy = (x,y) | |
| return img | |
| def draw_output_original(out, img): | |
| img = img.copy() | |
| for j in range(out['lf'][0].shape[0]): | |
| begin = out['beginning'][j] | |
| end = out['ending'][j] | |
| last_xy = None | |
| # for i in xrange(len(out['lf'])): | |
| begin_f = int(np.floor(begin)) | |
| end_f = int(np.ceil(end)) | |
| for i in range(begin_f, end_f+1): | |
| if i == begin_f: | |
| p0 = out['lf'][i][j].mean(axis=1) | |
| p1 = out['lf'][i+1][j].mean(axis=1) | |
| t = begin - np.floor(begin) | |
| p = p0 * (1 - t) + p1 * t | |
| elif i == end_f: | |
| p0 = out['lf'][i-1][j].mean(axis=1) | |
| if i != len(out['lf']): | |
| p1 = out['lf'][i][j].mean(axis=1) | |
| t = end - np.floor(end) | |
| p = p0 * (1 - t) + p1 * t | |
| else: | |
| p = p0 | |
| else: | |
| p = out['lf'][i][j].mean(axis=1) | |
| x = p[0] | |
| y = p[1] | |
| x = int(x) | |
| y = int(y) | |
| color = (0,0,0) | |
| cv2.circle(img,(x,y), 4, color, -1) | |
| if last_xy is not None: | |
| cv2.line(img, (x,y), last_xy, color, 2) | |
| last_xy = (x,y) | |
| for i in range(out['sol'].shape[0]): | |
| p = out['sol'][i] | |
| c = int(255 * p[-1]) | |
| color = (c,0,255-c) | |
| x = p[0] | |
| y = p[1] | |
| r = p[2] | |
| x_comp = np.cos(r) | |
| y_comp = -np.sin(r) | |
| s = p[3] | |
| rx = x + s * x_comp * 2 | |
| ry = y + s * y_comp * 2 | |
| rx2 = x - s * x_comp | |
| ry2 = y - s * y_comp | |
| rx = int(rx) | |
| ry = int(ry) | |
| rx2 = int(rx2) | |
| ry2 = int(ry2) | |
| x = int(x) | |
| y = int(y) | |
| scale = abs(int(s)) | |
| # color = (0,0,255) | |
| cv2.circle(img,(x,y), int(scale), color, 2) | |
| cv2.circle(img,(x,y), 4, color, -1) | |
| cv2.arrowedLine(img, (x,y), (rx,ry), color, 2, tipLength=0.25) | |
| # cv2.line(img, (rx2,ry2), (rx,ry), color, 2) | |
| cv2.putText(img,str(i),(x,y), cv2.FONT_HERSHEY_SIMPLEX, 1,(0,255,0),2,cv2.LINE_AA) | |
| return img | |