Merge pull request #477 from kagenihisomi/panel_order

Manga panel order finder

manga_translator/detection/panel_finder.py

@@ -0,0 +1,250 @@
+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])