123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250 |
- from pathlib import Path
- import sys
- import cv2 as cv
- import numpy as np
- from PIL import Image
- KERNEL_SIZE = 7
- BORDER_SIZE = 10
- def panel_process_image(img: Image.Image):
- """Preprocesses an image to make it easier to find panels.
- Args:
- img: The image to preprocess.
- Returns:
- The preprocessed image.
- """
- img_gray = cv.cvtColor(np.array(img), cv.COLOR_BGR2GRAY)
- img_gray = cv.GaussianBlur(img_gray, (KERNEL_SIZE, KERNEL_SIZE), 0)
- img_gray = cv.threshold(img_gray, 200, 255, cv.THRESH_BINARY)[1]
- # Add black border to image, to help with finding contours
- img_gray = cv.copyMakeBorder(
- img_gray,
- BORDER_SIZE,
- BORDER_SIZE,
- BORDER_SIZE,
- BORDER_SIZE,
- cv.BORDER_CONSTANT,
- value=255,
- )
- # Invert image
- img_gray = cv.bitwise_not(img_gray)
- return img_gray
- def remove_contained_contours(polygons):
- """Removes polygons from a list if any completely contain the other.
- Args:
- polygons: A list of polygons.
- Returns:
- A list of polygons with any contained polygons removed.
- """
- # Create a new list to store the filtered polygons.
- filtered_polygons = []
- # Iterate over the polygons.
- for polygon in polygons:
- # Check if the polygon contains any of the other polygons.
- contains = False
- for other_polygon in polygons:
- # Check if the polygon contains the other polygon and that the polygons
- if np.array_equal(other_polygon, polygon):
- continue
- rect1 = cv.boundingRect(other_polygon)
- rect2 = cv.boundingRect(polygon)
- # Check if rect2 is completely within rect1
- if (
- rect2[0] >= rect1[0]
- and rect2[1] >= rect1[1]
- and rect2[0] + rect2[2] <= rect1[0] + rect1[2]
- and rect2[1] + rect2[3] <= rect1[1] + rect1[3]
- ):
- contains = True
- break
- # If the polygon does not contain any of the other polygons, add it to the
- # filtered list.
- if not contains:
- filtered_polygons.append(polygon)
- return filtered_polygons
- def calc_panel_contours(im: Image.Image):
- img_gray = panel_process_image(im)
- contours_raw, hierarchy = cv.findContours(
- img_gray, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE
- )
- contours = contours_raw
- min_area = 10000
- contours = [i for i in contours if cv.contourArea(i) > min_area]
- contours = [cv.convexHull(i) for i in contours]
- contours = remove_contained_contours(contours)
- # Remap the contours to the original image
- contours = [i + np.array([[-BORDER_SIZE, -BORDER_SIZE]]) for i in contours]
- # Sort the contours by their y-coordinate.
- contours = order_panels(contours, img_gray)
- return contours
- def determine_panel_order_from_contours(contours):
- """
- build a tree of regions that are determined vertically
- order by an n like pattern
- """
- def draw_contours(im, contours):
- """Debugging, draws the contours on the image."""
- colors = [
- (255, 0, 0),
- (0, 255, 0),
- (0, 0, 255),
- ]
- im_contour = np.array(im)
- for i, contour in enumerate(range(len(contours))):
- color = colors[i % len(colors)]
- im_contour = cv.drawContours(im_contour, contours, i, color, 4, cv.LINE_AA)
- # Draw a number at the top left of contour
- x, y, _, _ = cv.boundingRect(contours[i])
- cv.putText(
- im_contour,
- str(i),
- (x + 50, y + 50),
- cv.FONT_HERSHEY_SIMPLEX,
- 1,
- color,
- 2,
- cv.LINE_AA,
- )
- img = Image.fromarray(im_contour)
- return img
- def save_draw_contours(pth: Path | str):
- if str:
- pth = Path(pth)
- pth_out = pth.parent / (pth.stem + "-contours")
- if not pth_out.exists():
- pth_out.mkdir()
- # Glob get all images in folder
- pths = [i for i in pth.iterdir() if i.suffix in [".png", ".jpg", ".jpeg"]]
- for t in pths:
- print(t)
- im = Image.open(t)
- contours = calc_panel_contours(im)
- img_panels = draw_contours(im, contours)
- f_name = t.stem + t.suffix
- img_panels.save(pth_out / f_name)
- def order_panels(contours, img_gray):
- """Orders the panels in a comic book page.
- Args:
- contours: A list of contours, where each contour is a list of points.
- Returns:
- A list of contours, where each contour is a list of points, ordered by
- their vertical position.
- """
- # Get the bounding boxes for each contour.
- bounding_boxes = [cv.boundingRect(contour) for contour in contours]
- # Generate groups of vertically overlapping bounding boxes.
- groups_indices = generate_vertical_bounding_box_groups_indices(bounding_boxes)
- c = []
- for group in groups_indices:
- # Reorder contours based on reverse z-order,
- cs = [bounding_boxes[i] for i in group]
- ymax, xmax = img_gray.shape
- order_scores = [1 * (ymax - i[1]) + i[0] * 1 for i in cs]
- # Sort the list based on the location score value
- combined_list = list(zip(group, order_scores))
- sorted_list = sorted(combined_list, key=lambda x: x[1], reverse=True)
- c.extend(sorted_list)
- ordered_contours = [contours[i[0]] for i in c]
- return ordered_contours
- def generate_vertical_bounding_box_groups_indices(bounding_boxes):
- """Generates groups of vertically overlapping bounding boxes.
- Args:
- bounding_boxes: A list of bounding boxes, where each bounding box is a tuple
- of (x, y, width, height).
- Returns:
- A list of groups, where each group is a list of bounding boxes that overlap
- vertically.
- """
- # Operate on indices Sort the bounding boxes by their y-coordinate.
- bbox_inds = np.argsort([i[1] for i in bounding_boxes])
- # generate groups of vertically overlapping bounding boxes
- groups = [[bbox_inds[0]]]
- for i in bbox_inds[1:]:
- is_old_group = False
- bbox = bounding_boxes[i]
- start1 = bbox[1]
- end1 = bbox[1] + bbox[3]
- for n, group in enumerate(groups):
- for ind in group:
- _bbox = bounding_boxes[ind]
- start2 = _bbox[1]
- end2 = _bbox[1] + _bbox[3]
- # Check for any partial overlapping
- if check_overlap((start1, end1), (start2, end2)):
- groups[n] = group + [i]
- is_old_group = True
- break
- if is_old_group:
- break
- else:
- groups.append([i])
- return groups
- def check_overlap(range1, range2):
- # Check if range1 is before range2
- if range1[1] < range2[0]:
- return False
- # Check if range1 is after range2
- elif range1[0] > range2[1]:
- return False
- # If neither of the above conditions are met, the ranges must overlap
- else:
- return True
- if __name__ == "__main__":
- save_draw_contours(sys.argv[1])
|