File size: 3,350 Bytes
c843d82 |
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 |
#!/usr/bin/env python3
"""
Developed by Nikhil Nageshwar Inturi
This module provides ImageSplitter for splitting PNG images into
equal-sized sub-images and saving them to an output directory.
"""
# imports
from pathlib import Path
import numpy as np, cv2, logging
from PIL import Image
# local imports
from utils.constants import setup_logging
Image.MAX_IMAGE_PIXELS = None
class ImageSplitter:
"""
Splits all PNG images in a source directory into sub-images of specified width and height.
"""
def __init__(self, source_dir: Path, output_dir: Path, sub_image_width: int, sub_image_height: int) -> None:
self.source_dir = Path(source_dir)
self.output_dir = Path(output_dir)
self.sub_w = sub_image_width
self.sub_h = sub_image_height
self.logger = logging.getLogger(self.__class__.__name__)
self._setup_output_directory()
def _setup_output_directory(self) -> None:
try:
self.output_dir.mkdir(parents=True, exist_ok=True)
self.logger.debug(f"Output directory ready: {self.output_dir}")
except Exception as e:
self.logger.error(f"Failed to create output directory {self.output_dir}: {e}")
raise
def split_all(self) -> None:
"""
Iterate over all PNG files in source_dir and split them.
"""
png_files = list(self.source_dir.glob("*.png"))
if not png_files:
self.logger.warning(f"No .png files found in {self.source_dir}")
return
for png_file in png_files:
try:
self.split_file(png_file)
except Exception:
self.logger.exception(f"Error splitting file: {png_file}")
def split_file(self, png_path: Path) -> None:
"""
Split a single PNG image into sub-images.
Args:
png_path: Path to the input .png file.
"""
with Image.open(png_path) as pil_img:
img = np.array(pil_img)
self.logger.debug(f"Loaded {png_path.name} with shape {img.shape}")
height, width = img.shape[:2]
cols = (width + self.sub_w - 1) // self.sub_w
rows = (height + self.sub_h - 1) // self.sub_h
for row in range(rows):
for col in range(cols):
x0 = col * self.sub_w
y0 = row * self.sub_h
x1 = min(x0 + self.sub_w, width)
y1 = min(y0 + self.sub_h, height)
sub_img = img[y0:y1, x0:x1]
output_name = f"{png_path.stem}_{row}_{col}.png"
output_path = self.output_dir / output_name
# success = cv2.imwrite(str(output_path), sub_img)
sub_img_bgr = cv2.cvtColor(sub_img, cv2.COLOR_RGB2BGR)
success = cv2.imwrite(str(output_path), sub_img_bgr)
if success:
self.logger.info(f"Saved sub-image: {output_name}")
else:
self.logger.error(f"Failed to save sub-image: {output_name}")
# testing
# from bin.constants import setup_logging
# from bin.generate_split_images import ImageSplitter
# setup_logging(logging.INFO)
# splitter = ImageSplitter(source_dir=Path("path/to/pngs"), output_dir=Path("path/to/splits"), sub_image_width=640, sub_image_height=640)
# splitter.split_all() |