markup: move functionality to library
This commit is contained in:
@ -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)
|
||||
|
||||
Reference in New Issue
Block a user