from typing import List, Tuple from pathlib import Path from .config import Config import numpy as np import cv2 class ImageProcessor: """Handles image preprocessing operations.""" def __init__(self, config: Config): self.config = config def mask_text_regions(self, input_path, bboxes: List[List[int]], output_filename: str = "1_text_removed.jpg", color: Tuple[int, int, int] = (0, 0, 0)) -> str: """Mask text regions in the image to reduce panel extraction noise.""" image = cv2.imread(input_path) if image is None: raise FileNotFoundError(f"Could not load image: {input_path}") for bbox in bboxes: x1, y1, x2, y2 = bbox cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness=-1) output_path = f'{self.config.output_folder}/{output_filename}' cv2.imwrite(output_path, image) print(f"✅ Text-masked image saved to: {output_path}") return str(output_path) def preprocess_image(self, processed_image_path) -> Tuple[str, str, str]: """Preprocess image for panel extraction.""" image = cv2.imread(processed_image_path) if image is None: raise FileNotFoundError(f"Could not load image: {processed_image_path}") # Convert to grayscale and binary gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) _, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV) is_inverted = False # binary, is_inverted = self.invert_if_black_dominates(binary) if not is_inverted: # Dilate to strengthen borders kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) dilated = cv2.dilate(binary, kernel, iterations=2) else: dilated = binary # Save intermediate results gray_path = f'{self.config.output_folder}/2_gray.jpg' binary_path = f'{self.config.output_folder}/3_binary.jpg' dilated_path = f'{self.config.output_folder}/4_dilated.jpg' cv2.imwrite(str(gray_path), gray) cv2.imwrite(str(binary_path), binary) cv2.imwrite(str(dilated_path), dilated) return str(gray_path), str(binary_path), str(dilated_path), is_inverted def invert_if_black_dominates(self, binary): # Threshold to binary image _, binary = cv2.threshold(binary, 127, 255, cv2.THRESH_BINARY) # Count black and white pixels black_pixels = np.sum(binary == 0) white_pixels = np.sum(binary == 255) # If black dominates, invert if black_pixels > white_pixels: print("🔄 Inverting image because black > white") inverted = cv2.bitwise_not(binary) else: print("✅ No inversion needed, white >= black") inverted = binary # Save result return inverted, black_pixels > white_pixels def remove_inner_sketch(self, input_path, output_filename="5_remove_inner_sketch.jpg", min_area_ratio=0): img = cv2.imread(input_path, cv2.IMREAD_GRAYSCALE) height, width = img.shape # Threshold image to binary _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV) # Find all contours contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Create mask for large contours (likely panel borders) mask = np.zeros_like(binary) for cnt in contours: area = cv2.contourArea(cnt) if area >= (height * width * min_area_ratio): cv2.drawContours(mask, [cnt], -1, 255, thickness=cv2.FILLED) # Apply mask to original image (keeps only large borders) cleaned = cv2.bitwise_and(binary, binary, mask=mask) # Optional: Apply morphological opening to clean tiny sketch lines kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel) # Invert back if needed cleaned = cv2.bitwise_not(cleaned) # Save output_path = f'{self.config.output_folder}/{output_filename}' cv2.imwrite(output_path, cleaned) print(f"✅ Remove Inner Sketch image saved to: {output_path}") return str(output_path) def thin_image_borders(self, processed_image_path: str, output_filename: str = "6_thin_border.jpg") -> str: """ Clean dilated image by thinning thick borders and removing hanging clusters. """ from skimage.morphology import skeletonize, remove_small_objects from skimage.measure import label # Load image img = cv2.imread(processed_image_path, cv2.IMREAD_GRAYSCALE) _, binary = cv2.threshold(img, 128, 1, cv2.THRESH_BINARY_INV) # invert, binary mask (0,1) # Skeletonize skeleton = skeletonize(binary).astype(np.uint8) # Remove small hanging clusters labeled = label(skeleton, connectivity=2) cleaned = remove_small_objects(labeled, min_size=150) # Adjust min_size for more/less pruning # Convert back to 0–255 uint8 image final = (cleaned > 0).astype(np.uint8) * 255 # Invert back if needed result = 255 - final # Save output_path = f'{self.config.output_folder}/{output_filename}' cv2.imwrite(output_path, result) print(f"✅ Cleaned and thinned image saved to: {output_path}") return str(output_path) def clean_dilated_image(self, dilated_path: str, output_filename: str = "6_dilated_cleaned.jpg", max_neighbors: int = 2) -> str: """Clean dilated image by thinning thick borders.""" dilated = cv2.imread(dilated_path, cv2.IMREAD_GRAYSCALE) if dilated is None: raise FileNotFoundError(f"Could not load dilated image: {dilated_path}") binary = (dilated == 0).astype(np.uint8) padded = np.pad(binary, pad_width=1, mode="constant", constant_values=0) cleaned = binary.copy() height, width = binary.shape row_black_counts = np.sum(binary, axis=1) for y in range(1, height + 1): for x in range(1, width + 1): if padded[y, x] == 1: neighbors = np.sum(padded[y-1:y+2, x-1:x+2]) - 1 if neighbors > max_neighbors: neighbor_rows = [r for r in [y-1, y, y+1] if 1 <= r <= height] if neighbor_rows: row_to_clear = min(neighbor_rows, key=lambda r: row_black_counts[r-1]) if y == row_to_clear: cleaned[y-1, x-1] = 0 cleaned_img = (1 - cleaned) * 255 output_path = f'{self.config.output_folder}/{output_filename}' cv2.imwrite(str(output_path), cleaned_img) print(f"✅ Cleaned dilated image saved to: {output_path}") return str(output_path)