109 lines
3.8 KiB
Python
109 lines
3.8 KiB
Python
import os
|
|
import shutil
|
|
|
|
from PIL import Image, ImageFilter, ImageDraw, ImageFont
|
|
import numpy
|
|
import imutils
|
|
import cv2
|
|
import dumper
|
|
import random as rng
|
|
from pathlib import Path
|
|
|
|
from .utils import cv2_rect, ensure_dir, set_file_perms, WORKDIR
|
|
|
|
# https://www.pyimagesearch.com/2014/10/20/finding-shapes-images-using-python-opencv/
|
|
|
|
|
|
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')
|
|
if not img.mode in ('RGBA', 'LA'):
|
|
print('no alpha channel: {}'.format(img.mode))
|
|
return None
|
|
|
|
alpha_layer = img.convert('RGBA').split()[-1]
|
|
|
|
alpha_layer = alpha_layer.filter(ImageFilter.GaussianBlur(5))
|
|
|
|
threshold = 5
|
|
alpha_layer = alpha_layer.point(lambda p: p > threshold and 255)
|
|
threshold = numpy.array(alpha_layer)
|
|
|
|
# alternate method
|
|
# blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
|
# thresh = cv2.threshold(blurred, 60, 255, cv2.THRESH_BINARY)[1]
|
|
|
|
thresh_path = str(path.with_suffix('.thresh.png'))
|
|
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)
|
|
|
|
bboxes = []
|
|
for c in contours:
|
|
# bounding rect
|
|
x, y, w, h = cv2.boundingRect(c)
|
|
# essentially center of mass
|
|
# NOT center of the bbox!
|
|
# M = cv2.moments(c)
|
|
# if M["m00"] == 0: M["m00"] = 0.00001
|
|
# cX = int(M["m10"] / M["m00"])
|
|
# cY = int(M["m01"] / M["m00"])
|
|
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)):
|
|
color = (rng.randint(0,512), rng.randint(0,512), rng.randint(0,512))
|
|
cv2.drawContours(contour_image, contours, i, color)
|
|
rect = bboxes[i]
|
|
cv2.rectangle(contour_image, (rect.left, rect.top), (rect.right, rect.bottom), color, 1)
|
|
# cv2.circle(contour_image, (cX, cY), 2, color, -1)
|
|
# cv2.putText(contour_image, "center", (cX - 20, cY - 15),
|
|
# cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
|
|
|
|
contour_path = str(path.with_suffix('.contour.png'))
|
|
cv2.imwrite(contour_path, contour_image)
|
|
os.chmod(contour_path, 0o664)
|
|
shutil.chown(contour_path, group='procat')
|
|
|
|
return img.width, img.height, bboxes
|
|
|
|
|
|
def write_debug_image(workdir, page_num, prods, scribbles):
|
|
"""Draw an image with boxes for products, images, and shapes."""
|
|
ensure_dir(workdir)
|
|
path = os.path.join(workdir, f"debug-page{page_num:03d}.png")
|
|
|
|
pagew = int(11*72)
|
|
pageh = int(8.5*72)
|
|
|
|
img = Image.new('RGBA', (pagew, pageh), 'white')
|
|
draw = ImageDraw.Draw(img, 'RGBA')
|
|
fnt = ImageFont.truetype('/usr/share/fonts/truetype/lato/Lato-Regular.ttf', 10)
|
|
|
|
for prod in filter(lambda p: p['page'] == page_num, prods):
|
|
rect = prod['rect']
|
|
fill_color = "hsv(120, 22%, 100%)" if 'matched' in prod else None
|
|
outline_color = "hsv(120, 50%, 100%)"
|
|
draw.rectangle((rect.p1(pageh), rect.p2(pageh)),
|
|
fill=fill_color, outline=outline_color, width=2)
|
|
bl = rect.p1(pageh)
|
|
draw.text((bl[0] + 3, bl[1] + 3), prod['material'],
|
|
font=fnt, fill="hsv(120, 22%, 50%)")
|
|
|
|
for scribble in filter(lambda s: s['page'] == page_num, scribbles):
|
|
rect = scribble['rect']
|
|
draw.rectangle((rect.p1(pageh), rect.p2(pageh)), outline="hsv(210, 22%, 100%)", width=2)
|
|
for box in scribble['bboxes']:
|
|
draw.rectangle((box.p1(pageh), box.p2(pageh)), outline="hsv(0, 22%, 100%)", width=2)
|
|
|
|
img.save(path)
|
|
set_file_perms(path)
|