Spaces:
Running
Running
File size: 3,517 Bytes
ad8cacf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
"""
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)
|