Unverified Commit e200b533 authored by tilman's avatar tilman
Browse files

add kmeans to main script

parent c9206758
......@@ -18,18 +18,21 @@ SHOW_WINDOW = True
SAVE_FILE = not SHOW_WINDOW
SAVE_BDCN = False
BDCN_EDGES = True
WHITE_OUTPUT_CANVAS = True
WHITE_OUTPUT_CANVAS = False
TRIANGLES = False
BODY_LINES = True
BISEC_VECTORS = False
BISEC_CONES = True
GLOBAL_LINE = True
INPAINT_AND_KMEANS = True
DRAW_FIRST_CONE_LAYER_BRIGTHER = True
BISEC_CONE_ANGLE = 50
CORRECTION_ANGLE = 23
OVERLAY_ALPHA = 0.2
BISEC_SKIP_LOWER_LEVELS = False
DISPLAY_RASTER_ELEMENTS = 500
KMEANS_AREA_MASK_THRESHOLD = 0.1 #max 1.0 (percent of pixels)
KMEANS_K = 5 #10
OUT_DIR = 'images/out/images_imdahl/main_structure_final_canvas'
#IN_DIR = "images/first_email/" # images from first email
IN_DIR = "images/images_imdahl/" # images from imdahl
......@@ -79,7 +82,7 @@ for img_name, img_bdcn in list(zip(images, images_bdcn))[15:16]: #jesus
# Process Image
print("calculating: "+img_name)
img = cv2.imread(img_name)
img_output_canvas = np.array([[[255,255,255]]*len(img[0])]*len(img),np.uint8) if WHITE_OUTPUT_CANVAS else cv2.imread(img_bdcn)
output_canvas = np.array([[[255,255,255]]*len(img[0])]*len(img),np.uint8) if WHITE_OUTPUT_CANVAS else cv2.imread(img_bdcn)
max_lw = max(len(img),len(img[0]))
esz = max_lw / DISPLAY_RASTER_ELEMENTS
if SKIP_OPENPOSE:
......@@ -92,9 +95,76 @@ for img_name, img_bdcn in list(zip(images, images_bdcn))[15:16]: #jesus
opWrapper.emplaceAndPop([datum])
img = datum.cvOutputData
if INPAINT_AND_KMEANS:
target = cv2.imread(img_name)
fposes = np.array([np.array([line[:2] for line in pose if line[2] > 0]) for pose in datum.poseKeypoints]) #filtered poses without zero lines
mask = np.zeros((len(img),len(img[0]),1), np.uint8)
kmout_mask = np.zeros((len(img),len(img[0]),1), np.uint8)
for pose in fposes:
convexhull = Polygon(pose).convex_hull
sconvexhull = affinity.scale(convexhull, xfact=1.6, yfact=1.6, origin=convexhull.centroid)
cv2.drawContours(mask, [polyToArr(sconvexhull)], 0, 255, int(10*esz))
cv2.drawContours(kmout_mask, [polyToArr(sconvexhull)], 0, 255, int(10*esz))
cv2.drawContours(mask, [polyToArr(sconvexhull)], 0, 255, -1)
cv2.drawContours(kmout_mask, [polyToArr(sconvexhull)], 0, 255, -1)
cv2.rectangle(mask, (0,0), (len(img[0]),len(img)), 255, int(40*esz))
inpainted = cv2.inpaint(target, mask, 3, cv2.INPAINT_TELEA)
inpainted = cv2.medianBlur(inpainted,int(esz*5)+1 if int(esz*5)%2==0 else int(esz*5))
print("kmeans_med1")
inpainted = cv2.bilateralFilter(inpainted,30,60,60)
print("kmeans_bil1")
# inpainted = cv2.bilateralFilter(inpainted,30,60,60)
# print("kmeans_bil2")
# inpainted = cv2.bilateralFilter(inpainted,30,60,60)
# print("kmeans_bil3")
# inpainted = cv2.bilateralFilter(inpainted,30,60,60)
Z = inpainted.reshape((-1,3))
Z = np.float32(Z)
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
K = KMEANS_K
ret, label, center = cv2.kmeans(Z,K,None,criteria,10,cv2.KMEANS_RANDOM_CENTERS)
center = np.uint8(center)
res = center[label.flatten()]
kmeans = res.reshape((inpainted.shape))
# cv2.imshow(img_name+"inpainted", inpainted)
# cv2.imshow(img_name+"mask", kmout_mask)
kmout_mask=cv2.cvtColor(kmout_mask,cv2.COLOR_GRAY2BGR)#change mask to a 3 channel image
km_mask_out=cv2.subtract(kmout_mask,kmeans)
km_mask_out=cv2.subtract(kmout_mask,km_mask_out)
#idea: count how many pixels the white mask has, count how many pixels each color from kmeans result appears in mask -> area > 80% is foreground
km_colors, km_counts = np.unique(km_mask_out.reshape(-1, km_mask_out.shape[-1]),axis=0,return_counts=True)
mask_colors, mask_counts = np.unique(kmout_mask.reshape(-1, kmout_mask.shape[-1]),axis=0,return_counts=True)
threshold_count = int(KMEANS_AREA_MASK_THRESHOLD*mask_counts[(np.argwhere(mask_colors>0))[0,0]]) #select counts from white color and generate threshold with it.
filtered_km_colors = km_colors[np.argwhere(km_counts>threshold_count)]
filtered_km_colors = filtered_km_colors[np.sum(filtered_km_colors, axis=2)>0] #filter out black from mask
print(km_colors, km_counts)
print(mask_colors, mask_counts)
print(filtered_km_colors)
print(len(km_colors))
cv2.putText(km_mask_out, "Foreground colors:", (10,25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255))
offset=220;
for color in filtered_km_colors:
print("color",color[0],color[1],color[2])
x = offset
y = 15
cv2.rectangle(km_mask_out, (x,y), (x+10,y+10), (int(color[0]),int(color[1]),int(color[2])), -1)
offset += 15
cv2.imshow(img_name+"kmeans_masked", km_mask_out)
#replace colors
cv2.imshow(img_name+"kmeans", kmeans)
cv2.waitKey(0)
if BDCN_EDGES:
overlay = np.array([[[255,255,255]]*len(img[0])]*len(img),np.uint8) #white image canvas
output_canvas = overlay_two_image_v2(img_output_canvas, overlay, [0,0,0], 0.2)
#overlay = np.array([[[255,255,255]]*len(img[0])]*len(img),np.uint8) #white image canvas
#output_canvas = overlay_two_image_v2(output_canvas, overlay, [0,0,0], 0.2)
fposes = np.array([np.array([line[:2] for line in pose if line[2] > 0]) for pose in datum.poseKeypoints]) #filtered poses without zero lines
for pose in fposes:
convexhull = Polygon(pose).convex_hull
......@@ -109,6 +179,17 @@ for img_name, img_bdcn in list(zip(images, images_bdcn))[15:16]: #jesus
# cv2.rectangle(img, (0,0), (len(img[0]),len(img)), (255,0,0), int(40*esz))
cv2.rectangle(output_canvas, (0,0), (len(img[0]),len(img)), (255,255,255), int(40*esz))
output_canvas = cv2.medianBlur(output_canvas,int(esz*5)+1 if int(esz*5)%2==0 else int(esz*5))
print("med1")
output_canvas = cv2.bilateralFilter(output_canvas,30,60,60)
print("bil1")
output_canvas = cv2.bilateralFilter(output_canvas,30,60,60)
print("bil2")
output_canvas = cv2.bilateralFilter(output_canvas,30,60,60)
print("bil3")
output_canvas = cv2.bilateralFilter(output_canvas,30,60,60)
print("bil4")
if TRIANGLES or BODY_LINES:
triangles = [poseToTriangle(pose) for pose in datum.poseKeypoints]
for triangle in triangles:
......@@ -166,7 +247,7 @@ for img_name, img_bdcn in list(zip(images, images_bdcn))[15:16]: #jesus
if SAVE_FILE:
cv2.imwrite(os.path.basename(img_name),img)
if BDCN_EDGES:
cv2.imwrite('CANVAS_'+os.path.basename(img_output_canvas),output_canvas)
cv2.imwrite('CANVAS_'+os.path.basename(img_bdcn),output_canvas)
if SHOW_WINDOW:
cv2.namedWindow(img_name, cv2.WINDOW_NORMAL)
cv2.imshow(img_name, img)
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment