Spaces:
Running
Running
| """ | |
| Image refinement functionality. | |
| """ | |
| import os | |
| import cv2 | |
| import numpy as np | |
| from loguru import logger | |
| from .splitter import create_output_directory | |
| def refine_animal_illustrations(input_dir='split_animals', refined_dir='refined_animals'): | |
| """ | |
| Refine the extracted animal illustrations by removing excess whitespace. | |
| Args: | |
| input_dir (str): Directory containing the initially split animal images | |
| refined_dir (str): Directory where refined animal images will be saved | |
| Returns: | |
| list: Paths to the refined animal images | |
| """ | |
| # Create refined output directory | |
| refined_dir = create_output_directory(refined_dir) | |
| refined_paths = [] | |
| # Process each file in the input directory | |
| for filename in os.listdir(input_dir): | |
| if not filename.lower().endswith(('.png', '.jpg', '.jpeg')): | |
| continue | |
| file_path = os.path.join(input_dir, filename) | |
| try: | |
| # Read image | |
| img = cv2.imread(file_path) | |
| if img is None: | |
| logger.warning(f"Could not read image: {file_path}") | |
| continue | |
| # Process image | |
| refined_img = remove_excess_whitespace(img) | |
| # Save refined image | |
| refined_path = os.path.join(refined_dir, filename) | |
| cv2.imwrite(refined_path, refined_img) | |
| refined_paths.append(refined_path) | |
| logger.info(f"Refined: {refined_path}") | |
| except Exception as e: | |
| logger.error(f"Error processing {file_path}: {str(e)}") | |
| continue | |
| logger.info(f"Refinement complete. Images saved to {refined_dir}") | |
| return refined_paths | |
| def remove_excess_whitespace(img, padding=10): | |
| """ | |
| Remove excess whitespace around the main content of an image. | |
| Args: | |
| img (numpy.ndarray): Input image | |
| padding (int): Padding to add around content | |
| Returns: | |
| numpy.ndarray: Cropped image | |
| """ | |
| # Convert to grayscale | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| # Apply threshold to separate content from background | |
| _, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV) | |
| # Find contours | |
| contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if not contours: | |
| return img | |
| # Find largest contour (assumed to be the main content) | |
| largest_contour = max(contours, key=cv2.contourArea) | |
| # Get bounding rectangle | |
| x, y, w, h = cv2.boundingRect(largest_contour) | |
| # Add padding | |
| x = max(0, x - padding) | |
| y = max(0, y - padding) | |
| w = min(img.shape[1] - x, w + 2 * padding) | |
| h = min(img.shape[0] - y, h + 2 * padding) | |
| # Crop image | |
| return img[y:y+h, x:x+w] | |
| def calculate_content_bounds(gray_img, threshold=240): | |
| """ | |
| Calculate content boundaries in grayscale image. | |
| Args: | |
| gray_img (numpy.ndarray): Grayscale image | |
| threshold (int): Brightness threshold | |
| Returns: | |
| tuple: (min_x, min_y, max_x, max_y) | |
| """ | |
| # Create binary mask | |
| _, mask = cv2.threshold(gray_img, threshold, 255, cv2.THRESH_BINARY_INV) | |
| # Find non-zero points | |
| points = cv2.findNonZero(mask) | |
| if points is None: | |
| return None | |
| # Calculate bounds | |
| x, y, w, h = cv2.boundingRect(points) | |
| return (x, y, x + w, y + h) | |