import logging import shutil from pathlib import Path from typing import Optional, Tuple import cv2 import numpy as np from tqdm import tqdm class MNISTProcessor: def __init__(self): self.logger = logging.getLogger(__name__) self.project_root = Path(__file__).parent.parent.parent.parent self.source_path = self.project_root / "data" / "raw" / "mnist" self.temp_path = self.project_root / "data" / "temp" / "MNIST" self.target_size = (27, 27) def process(self) -> bool: try: self.logger.info("Processing MNIST dataset...") self.temp_path.mkdir(parents=True, exist_ok=True) if not self.source_path.exists(): self.logger.error(f"Source dataset not found at {self.source_path}") return False self._copy_dataset() self._normalize_documentation() image_files = list((self.temp_path / "images").glob("*.png")) total_files = len(image_files) for idx, img_path in enumerate(tqdm(image_files, desc="Processing MNIST")): try: image = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE) if image is None: self.logger.warning(f"Failed to read image: {img_path}") continue processed = self._process_single_image(image) cv2.imwrite(str(img_path), processed) progress = (idx + 1) / total_files * 100 self.logger.debug(f"Processed {img_path.name} ({progress:.2f}%)") except Exception as e: self.logger.error(f"Error processing {img_path}: {str(e)}") continue self.logger.info("MNIST processing completed successfully") return True except Exception as e: self.logger.error(f"Error in MNIST processing: {str(e)}") return False def _process_single_image(self, image: np.ndarray) -> np.ndarray: """Process a single MNIST image through enhancement pipeline.""" # Enhance contrast image = self._enhance_contrast(image) # Remove noise image = self._reduce_noise(image) # Center character image = self._center_character(image) # Normalize stroke width image = self._normalize_stroke(image) # Resize to target size image = cv2.resize(image, self.target_size, interpolation=cv2.INTER_AREA) return image def _enhance_contrast(self, image: np.ndarray) -> np.ndarray: """Enhance image contrast using CLAHE.""" clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) return clahe.apply(image) def _reduce_noise(self, image: np.ndarray) -> np.ndarray: """Reduce noise while preserving edges.""" denoised = cv2.bilateralFilter(image, d=5, sigmaColor=75, sigmaSpace=75) _, binary = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) return binary def _center_character(self, image: np.ndarray) -> np.ndarray: """Center the character using moments.""" contours, _ = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: return image # Get largest contour largest_contour = max(contours, key=cv2.contourArea) # Calculate moments M = cv2.moments(largest_contour) if M["m00"] == 0: return image # Calculate center cx = int(M["m10"] / M["m00"]) cy = int(M["m01"] / M["m00"]) # Calculate shift rows, cols = image.shape shift_x = (cols // 2) - cx shift_y = (rows // 2) - cy # Create transformation matrix M = np.float32([[1, 0, shift_x], [0, 1, shift_y]]) # Apply shift return cv2.warpAffine(image, M, (cols, rows)) def _normalize_stroke(self, image: np.ndarray) -> np.ndarray: """Normalize stroke width using morphological operations.""" kernel_size = max(2, min(image.shape) // 30) kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) return cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel) def _normalize_documentation(self) -> None: """Normalize documentation to simplified format (filename label).""" doc_path = self.temp_path / "documentation.txt" normalized_path = self.temp_path / "normalized_documentation.txt" skip_lines = 4 # Number of header lines to skip with open(doc_path, 'r') as f, open(normalized_path, 'w') as out: # Skip header lines for _ in range(skip_lines): next(f, None) for line in f: if not line.strip() or line.startswith('#'): continue parts = line.strip().split() if len(parts) >= 2: filename = parts[0] # First column (filename) label = parts[-1] # Last column (label) out.write(f"{filename} {label}\n") normalized_path.replace(doc_path) self.logger.info("Documentation normalized to simplified format") def _copy_dataset(self) -> None: """Copy dataset files to temp directory.""" if self.source_path.exists(): shutil.copytree(self.source_path, self.temp_path, dirs_exist_ok=True) self.logger.info(f"Copied dataset to {self.temp_path}") else: raise FileNotFoundError(f"Source dataset not found at {self.source_path}") def test_processor(): """Test MNIST processor functionality.""" import tempfile test_dir = Path(tempfile.mkdtemp()) try: # Create test directory structure img_dir = test_dir / "data" / "raw" / "mnist" / "images" img_dir.mkdir(parents=True) # Create test image test_image = np.zeros((28, 28), dtype=np.uint8) cv2.rectangle(test_image, (8, 8), (20, 20), 255, -1) # Add noise noise = np.random.normal(0, 25, test_image.shape).astype(np.uint8) noisy_image = cv2.add(test_image, noise) # Save test image and documentation cv2.imwrite(str(img_dir / "test.png"), noisy_image) with open(test_dir / "data" / "raw" / "mnist" / "documentation.txt", 'w') as f: f.write("test.png 5\n") # Initialize processor processor = MNISTProcessor() processor.project_root = test_dir processor.source_path = test_dir / "data" / "raw" / "mnist" processor.temp_path = test_dir / "data" / "temp" / "MNIST" # Process image assert processor.process() == True # Verify processed image processed_path = processor.temp_path / "images" / "test.png" assert processed_path.exists() processed_img = cv2.imread(str(processed_path), cv2.IMREAD_GRAYSCALE) assert processed_img is not None assert processed_img.shape == (27, 27) assert np.mean(processed_img) > 0 assert len(np.unique(processed_img)) > 1 # Verify documentation with open(processor.temp_path / "documentation.txt", 'r') as f: content = f.read() assert "MNIST" in content assert "27 27" in content print("All tests passed!") finally: shutil.rmtree(test_dir) if __name__ == "__main__": logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) test_processor()