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)