| import cv2 | |
| import numpy as np | |
| import os | |
| from pathlib import Path | |
| from PIL import Image | |
| class AutomatedCollageSplitter: | |
| def __init__(self): | |
| self.min_segment_area_ratio = 0.01 | |
| self.max_segment_area_ratio = 0.95 | |
| self.min_aspect_ratio = 0.2 | |
| self.max_aspect_ratio = 5.0 | |
| self.min_solidity = 0.9 | |
| self.nms_threshold = 0.3 | |
| self.upscaler = None | |
| self._initialize_upscaler() | |
| def _initialize_upscaler(self): | |
| if self.upscaler is not None: | |
| return | |
| try: | |
| model_name = 'fsrcnn' | |
| model_path = os.path.join('models', 'weights', 'FSRCNN-small_x4.pb') | |
| scale = 4 | |
| if not os.path.isfile(model_path): | |
| raise FileNotFoundError(f"Model file not found at {model_path}") | |
| self.upscaler = cv2.dnn_superres.DnnSuperResImpl_create() | |
| self.upscaler.readModel(model_path) | |
| self.upscaler.setModel(model_name, scale) | |
| print(f"✓ OpenCV DNN upscaler model loaded: {model_name} with scale x{scale}") | |
| except Exception as e: | |
| print(f"⚠ Could not initialize OpenCV DNN upscaler: {e}. Proceeding without upscaling.") | |
| self.upscaler = None | |
| def _upscale_image(self, image_array): | |
| if not self.upscaler: | |
| print("❌ Upscaling skipped because the upscaler is not available.") | |
| return image_array | |
| try: | |
| return self.upscaler.upsample(image_array) | |
| except Exception as e: | |
| print(f"❌ Error during image upscaling: {e}") | |
| return image_array | |
| def preprocess_for_contours(self, image): | |
| """Pre-process the image to make panel borders stand out for contour detection.""" | |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| binary = cv2.adaptiveThreshold( | |
| gray, 255, | |
| cv2.ADAPTIVE_THRESH_MEAN_C, | |
| cv2.THRESH_BINARY_INV, | |
| 21, | |
| 8 | |
| ) | |
| kernel = np.ones((5, 5), np.uint8) | |
| closed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=2) | |
| return closed | |
| def find_panel_contours(self, processed_image, original_shape): | |
| """Find and filter contours that are likely to be comic panels.""" | |
| contours, _ = cv2.findContours(processed_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| height, width = original_shape[:2] | |
| total_area = height * width | |
| min_area = total_area * self.min_segment_area_ratio | |
| max_area = total_area * self.max_segment_area_ratio | |
| potential_panels = [] | |
| for contour in contours: | |
| area = cv2.contourArea(contour) | |
| if not (min_area < area < max_area): | |
| continue | |
| x, y, w, h = cv2.boundingRect(contour) | |
| if h == 0: continue | |
| aspect_ratio = w / h | |
| if not (self.min_aspect_ratio < aspect_ratio < self.max_aspect_ratio): | |
| continue | |
| hull = cv2.convexHull(contour) | |
| hull_area = cv2.contourArea(hull) | |
| if hull_area == 0: continue | |
| solidity = float(area) / hull_area | |
| if solidity < self.min_solidity: | |
| continue | |
| potential_panels.append([x, y, x + w, y + h, area]) | |
| return np.array(potential_panels) | |
| def apply_non_maximum_suppression(self, boxes): | |
| """Apply NMS to merge overlapping bounding boxes.""" | |
| if len(boxes) == 0: | |
| return [] | |
| boxes = boxes[boxes[:, 4].argsort()[::-1]] | |
| picked_boxes = [] | |
| while len(boxes) > 0: | |
| best_box = boxes[0] | |
| picked_boxes.append(best_box) | |
| remaining_boxes = boxes[1:] | |
| x1 = np.maximum(best_box[0], remaining_boxes[:, 0]) | |
| y1 = np.maximum(best_box[1], remaining_boxes[:, 1]) | |
| x2 = np.minimum(best_box[2], remaining_boxes[:, 2]) | |
| y2 = np.minimum(best_box[3], remaining_boxes[:, 3]) | |
| inter_w = np.maximum(0, x2 - x1) | |
| inter_h = np.maximum(0, y2 - y1) | |
| intersection_area = inter_w * inter_h | |
| best_box_area = (best_box[2] - best_box[0]) * (best_box[3] - best_box[1]) | |
| remaining_boxes_area = (remaining_boxes[:, 2] - remaining_boxes[:, 0]) * (remaining_boxes[:, 3] - remaining_boxes[:, 1]) | |
| union_area = best_box_area + remaining_boxes_area - intersection_area | |
| iou = intersection_area / union_area | |
| boxes = remaining_boxes[iou < self.nms_threshold] | |
| return np.array(picked_boxes) | |
| def split_collage(self, image_path, output_dir=None, debug=False): | |
| """Main function to automatically split collage using contour detection and NMS.""" | |
| img = cv2.imread(str(image_path)) | |
| if img is None: | |
| raise ValueError(f"Could not load image from {image_path}") | |
| print(f"Processing image: {image_path}") | |
| print(f"Image dimensions: {img.shape[1]}x{img.shape[0]}") | |
| processed_image = self.preprocess_for_contours(img) | |
| print("✓ Preprocessed image for contour detection") | |
| potential_panels = self.find_panel_contours(processed_image, img.shape) | |
| print(f"✓ Found {len(potential_panels)} potential panel contours") | |
| final_panels = self.apply_non_maximum_suppression(potential_panels) | |
| print(f"✓ Refined to {len(final_panels)} panels after Non-Maximum Suppression") | |
| if len(final_panels) > 0: | |
| img_height = img.shape[0] | |
| panel_heights = [box[3] - box[1] for box in final_panels] | |
| if panel_heights: | |
| max_panel_height = max(panel_heights) | |
| if max_panel_height > 0: | |
| bottom_margin = 10 | |
| height_threshold_ratio = 0.8 | |
| truly_final_panels = [] | |
| for box in final_panels: | |
| h = box[3] - box[1] | |
| y2 = box[3] | |
| is_at_bottom = y2 >= (img_height - bottom_margin) | |
| is_too_short = h < (max_panel_height * height_threshold_ratio) | |
| if is_at_bottom and is_too_short: | |
| print(f"Skipping potentially incomplete panel at the bottom (h={h} vs max_h={max_panel_height})") | |
| continue | |
| truly_final_panels.append(box) | |
| final_panels = truly_final_panels | |
| if len(final_panels) < 4: | |
| print("⚠ Contour detection found too few panels. Creating fallback 2x2 grid...") | |
| h, w = img.shape[:2] | |
| final_panels = np.array([ | |
| [0, 0, w//2, h//2, 0], | |
| [w//2, 0, w, h//2, 0], | |
| [0, h//2, w//2, h, 0], | |
| [w//2, h//2, w, h, 0] | |
| ]) | |
| final_panels = sorted(final_panels, key=lambda b: (b[1], b[0])) | |
| if output_dir is None: | |
| output_dir = Path(image_path).parent / f"{Path(image_path).stem}_segments" | |
| output_dir = Path(output_dir) | |
| output_dir.mkdir(exist_ok=True) | |
| saved_segments_info = [] | |
| for i, box in enumerate(final_panels): | |
| x1, y1, x2, y2, _ = map(int, box) | |
| padding = 3 | |
| x1 = max(0, x1 - padding) | |
| y1 = max(0, y1 - padding) | |
| x2 = min(img.shape[1], x2 + padding) | |
| y2 = min(img.shape[0], y2 + padding) | |
| segment = img[y1:y2, x1:x2] | |
| original_dims = (segment.shape[1], segment.shape[0]) | |
| upscaled_segment_np = self._upscale_image(segment) | |
| final_image = Image.fromarray(cv2.cvtColor(upscaled_segment_np, cv2.COLOR_BGR2RGB)) | |
| final_dims = (final_image.width, final_image.height) | |
| output_path = output_dir / f"segment_{i+1:02d}.png" | |
| final_image.save(str(output_path)) | |
| caption = ( | |
| f"Panel {i+1}<br>" | |
| f"Original: {original_dims[0]}x{original_dims[1]}<br>" | |
| f"Upscaled: {final_dims[0]}x{final_dims[1]}" | |
| ) | |
| saved_segments_info.append({ | |
| "path": str(output_path), | |
| "caption": caption | |
| }) | |
| print(f" Saved segment {i+1}: {final_dims[0]}x{final_dims[1]} pixels to {output_path}") | |
| if debug: | |
| debug_potential_panels = self.find_panel_contours(processed_image, img.shape) | |
| self.create_debug_images(img, processed_image, debug_potential_panels, final_panels, output_dir) | |
| print(f"\n🎉 Successfully split collage into {len(saved_segments_info)} segments!") | |
| print(f"📁 Segments saved in: {output_dir}") | |
| return saved_segments_info | |
| def create_debug_images(self, original, processed, potential_boxes, final_boxes, output_dir): | |
| """Create debug images showing the processing steps.""" | |
| cv2.imwrite(str(output_dir / "debug_01_binary_closed.png"), processed) | |
| potential_img = original.copy() | |
| if len(potential_boxes) > 0: | |
| for x1, y1, x2, y2, _ in potential_boxes: | |
| cv2.rectangle(potential_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 255), 2) | |
| cv2.imwrite(str(output_dir / "debug_02_potential_boxes.png"), potential_img) | |
| final_img = original.copy() | |
| if len(final_boxes) > 0: | |
| for x1, y1, x2, y2, _ in final_boxes: | |
| cv2.rectangle(final_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 3) | |
| cv2.imwrite(str(output_dir / "debug_03_final_panels.png"), final_img) | |
| print("🔍 Debug images saved:") | |
| print(" - debug_01_binary_closed.png (preprocessed)") | |
| print(" - debug_02_potential_boxes.png (before NMS)") | |
| print(" - debug_03_final_panels.png (after NMS)") | |
| def main(): | |
| """Example usage""" | |
| splitter = AutomatedCollageSplitter() | |
| image_path = "path/to/your/comic_image.png" | |
| try: | |
| if not Path(image_path).exists(): | |
| print(f"❌ Image file not found: {image_path}") | |
| print("Please update the image_path variable with the correct path to your collage image.") | |
| return | |
| segments = splitter.split_collage( | |
| image_path=image_path, | |
| debug=True | |
| ) | |
| print(f"\n📊 Processing complete!") | |
| print(f"Generated {len(segments)} separate images from the collage") | |
| except Exception as e: | |
| print(f"❌ Error processing image: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| if __name__ == "__main__": | |
| main() | |