markup: move functionality to library

This commit is contained in:
2019-10-18 15:14:05 -07:00
parent 5282f7cb2f
commit 94c1a419dc
5 changed files with 104 additions and 24 deletions

View File

@ -11,7 +11,7 @@ from pathlib import Path
from django.conf import settings
from .utils import cv2_rect
from .utils import cv2_rect, set_file_perms
WORKDIR = os.path.join(settings.ASSET_DIR, 'markup', 'work')
@ -19,6 +19,9 @@ WORKDIR = os.path.join(settings.ASSET_DIR, 'markup', 'work')
def find_shapes(image_path):
"""Find shapes in the image, returning bounding boxes around each.
Writes debug images next to the input image.
"""
path = Path(image_path)
img = Image.open(image_path, 'r')
@ -39,14 +42,12 @@ def find_shapes(image_path):
# thresh = cv2.threshold(blurred, 60, 255, cv2.THRESH_BINARY)[1]
thresh_path = str(path.with_suffix('.thresh.png'))
# print('write to', thresh_path)
cv2.imwrite(thresh_path, threshold)
os.chmod(thresh_path, 0o664)
shutil.chown(thresh_path, group='procat')
contours = cv2.findContours(threshold, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours = imutils.grab_contours(contours)
# print("{} shapes".format(len(contours)))
bboxes = []
for c in contours:
@ -58,13 +59,11 @@ def find_shapes(image_path):
# if M["m00"] == 0: M["m00"] = 0.00001
# cX = int(M["m10"] / M["m00"])
# cY = int(M["m01"] / M["m00"])
#print('add contour rect: {}'.format(cv2_rect(x, y, w, h)))
bboxes.append(cv2_rect(x, y, w, h))
# draw contours
contour_image = numpy.zeros((threshold.shape[0], threshold.shape[1], 3), dtype=numpy.uint8)
for i in range(len(contours)):
# compute the center of the contour
color = (rng.randint(0,512), rng.randint(0,512), rng.randint(0,512))
cv2.drawContours(contour_image, contours, i, color)
rect = bboxes[i]
@ -74,7 +73,6 @@ def find_shapes(image_path):
# cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
contour_path = str(path.with_suffix('.contour.png'))
#print('write to', contour_path)
cv2.imwrite(contour_path, contour_image)
os.chmod(contour_path, 0o664)
shutil.chown(contour_path, group='procat')
@ -82,8 +80,9 @@ def find_shapes(image_path):
return img.width, img.height, bboxes
def write_debug_image(cat_name, page_num, prods, scribbles):
path = os.path.join(WORKDIR, "debug-{}-{}.png".format(cat_name, page_num))
def write_debug_image(subdir, cat_name, page_num, prods, scribbles):
"""Draw an image with boxes for products, images, and shapes."""
path = os.path.join(WORKDIR, subdir, f"{cat_name}-debug-page{page_num:03d}.png")
pagew = int(11*72)
pageh = int(8.5*72)
@ -109,3 +108,4 @@ def write_debug_image(cat_name, page_num, prods, scribbles):
draw.rectangle((box.p1(pageh), box.p2(pageh)), outline="hsv(0, 22%, 100%)", width=2)
img.save(path)
set_file_perms(path)