commit 134eefe58dda8d12d22cf050368df7dda78e71a0 Author: HarishCM Date: Sat Jan 31 16:47:49 2026 +0530 Add files via upload diff --git a/README.md b/README.md new file mode 100644 index 0000000..9109b2f --- /dev/null +++ b/README.md @@ -0,0 +1,88 @@ +# ComfyUI Document Scanner + +A custom ComfyUI node that implements advanced document scanning with automatic perspective correction and multiple enhancement options. + +## Features + +- **Automatic Document Detection**: Uses contour detection to identify document boundaries +- **Perspective Correction**: Corrects skewed documents to flat, rectangular format +- **Multiple Enhancement Methods**: 6 different enhancement algorithms for optimal results +- **Batch Processing**: Handles multiple images at once +- **Debug Visualization**: Optional edge detection visualization +- **Fallback Safety**: Graceful handling of edge cases and errors + +## Nodes + +### Document Scanner +Full-featured node with all configuration options: + +**Inputs:** +- `image`: Input document image(s) +- `enhancement_method`: Choose from 6 enhancement methods + - `sharpening`: General purpose sharpening with HSV adjustments + - `cartooning`: Cartoon-like effect with edge detection + - `clahe`: Contrast Limited Adaptive Histogram Equalization + - `threshold`: Binary thresholding with Otsu's method + - `adaptive_threshold`: Local adaptive thresholding + - `flat_field`: Flat field correction for uneven lighting +- `edge_threshold_low/high`: Canny edge detection thresholds (20, 70 default) +- `blur_kernel_size`: Bilateral filter kernel size (5 default) +- `skip_preprocessing`: Skip GrabCut text removal step +- `return_debug_edges`: Output edge detection visualization + +**Outputs:** +- `scanned_image`: Final processed document +- `debug_edges`: Edge detection visualization (if enabled) + +### Simple Document Scanner +Simplified interface with preset configurations: + +**Inputs:** +- `image`: Input document image(s) +- `enhancement`: Choose enhancement preset + - `auto`: General purpose (sharpening) + - `text`: Optimized for text documents (adaptive threshold) + - `photo`: Optimized for photo documents (CLAHE) + +**Outputs:** +- `scanned_image`: Final processed document + +## Algorithm Overview + +1. **Preprocessing**: Optional GrabCut segmentation to remove text +2. **Grayscale Conversion**: Convert to single channel for processing +3. **Blur**: Bilateral filtering for edge-preserving noise reduction +4. **Edge Detection**: Canny edge detection to find document boundaries +5. **Contour Detection**: Find largest quadrilateral contour +6. **Perspective Correction**: Warp document to rectangular format +7. **Enhancement**: Apply selected enhancement method + +## Installation + +1. Copy the `comfyui_document_scanner` folder to your ComfyUI `custom_nodes` directory +2. Install dependencies: `pip install -r requirements.txt` +3. Restart ComfyUI +4. The nodes will appear under `image/processing` category + +## Requirements + +- opencv-python >= 4.5.0 +- numpy >= 1.20.0 +- torch >= 1.9.0 (provided by ComfyUI) + +## Error Handling + +The node includes robust error handling: +- If document detection fails, returns original image +- If perspective correction fails, skips correction step +- Graceful fallbacks for all processing steps +- Error messages logged to console + +## Tips for Best Results + +- Use well-lit images with clear document boundaries +- Ensure document occupies significant portion of image +- For text documents, try `adaptive_threshold` enhancement +- For photos/mixed content, try `clahe` or `sharpening` +- Adjust edge detection thresholds if having detection issues +- Use debug edges output to troubleshoot detection problems \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..1e97ced --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +from .document_scanner import DocumentScannerNode, SimpleDocumentScannerNode +from .black_bg_scanner import BlackBackgroundScannerNode + +NODE_CLASS_MAPPINGS = { + "DocumentScanner": DocumentScannerNode, + "SimpleDocumentScanner": SimpleDocumentScannerNode, + "BlackBackgroundScanner": BlackBackgroundScannerNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DocumentScanner": "Document Scanner", + "SimpleDocumentScanner": "Simple Document Scanner", + "BlackBackgroundScanner": "Black Background Scanner" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/black_bg_scanner.py b/black_bg_scanner.py new file mode 100644 index 0000000..e441500 --- /dev/null +++ b/black_bg_scanner.py @@ -0,0 +1,150 @@ +import cv2 +import numpy as np +import torch +from .utils import tensor_to_cv2, cv2_to_tensor, reorder, crop_out, enhance_image + + +def detect_object_on_black_background(image, threshold=30): + """Detect object boundaries against black background""" + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + + # Create binary mask: black background vs object + _, binary_mask = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) + + # Clean up small noise + kernel = np.ones((5, 5), np.uint8) + binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel, iterations=2) + binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel, iterations=1) + + return binary_mask + + +def find_object_contour_simple(binary_mask): + """Find main object contour from binary mask""" + contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if not contours: + # Fallback: return full image bounds + h, w = binary_mask.shape + return np.array([[0, 0], [w-1, 0], [w-1, h-1], [0, h-1]]) + + # Get largest contour (main object) + largest_contour = max(contours, key=cv2.contourArea) + + # Try to approximate to quadrilateral + perimeter = cv2.arcLength(largest_contour, True) + + # Start with loose approximation and tighten if needed + for epsilon_factor in [0.01, 0.02, 0.03, 0.05]: + approx = cv2.approxPolyDP(largest_contour, epsilon_factor * perimeter, True) + if len(approx) == 4: + return approx.reshape((4, 2)) + + # If quad approximation fails, use minimum area rectangle + rect = cv2.minAreaRect(largest_contour) + box = cv2.boxPoints(rect) + return np.int0(box) + + +def scan_black_background_object(image, enhancement_method="clahe", bg_threshold=30): + """ + Optimized scanning for objects on black background + + Args: + image: Input image (OpenCV format) + enhancement_method: Enhancement to apply + bg_threshold: Threshold to separate object from black background + """ + try: + # Step 1: Detect object against black background + binary_mask = detect_object_on_black_background(image, bg_threshold) + + # Step 2: Find object contour + vertices = find_object_contour_simple(binary_mask) + + # Step 3: Perspective correction + cropped = crop_out(image, vertices) + + # Step 4: Enhancement (preserve patterns/textures) + enhanced = enhance_image(cropped, enhancement_method) + + return enhanced, binary_mask + + except Exception as e: + print(f"Black background scanning error: {str(e)}") + return image, np.zeros_like(image[:,:,0]) + + +class BlackBackgroundScannerNode: + """ + Optimized ComfyUI node for objects on black backgrounds + """ + + CATEGORY = "image/processing" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "enhancement": (["clahe", "sharpening", "flat_field", "none"], { + "default": "clahe" + }), + "background_threshold": ("INT", { + "default": 30, + "min": 10, + "max": 100, + "step": 5 + }), + "return_mask": ("BOOLEAN", { + "default": False + }) + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("scanned_image", "detection_mask") + FUNCTION = "scan_black_background" + + def scan_black_background(self, image, enhancement, background_threshold, return_mask): + """ + Main function for black background object scanning + """ + try: + results = [] + masks = [] + + # Process each image in batch + for i in range(image.shape[0]): + # Convert to OpenCV format + cv2_image = tensor_to_cv2(image[i:i+1]) + + # Process with optimized algorithm + processed, mask = scan_black_background_object( + cv2_image, enhancement, background_threshold + ) + + # Convert back to tensors + result_tensor = cv2_to_tensor(processed) + results.append(result_tensor) + + if return_mask: + # Convert mask to 3-channel for visualization + mask_3ch = np.stack([mask, mask, mask], axis=2) + mask_tensor = cv2_to_tensor(mask_3ch) + masks.append(mask_tensor) + + # Combine results + final_result = torch.cat(results, dim=0) + + if return_mask and masks: + mask_result = torch.cat(masks, dim=0) + else: + mask_result = torch.zeros_like(final_result) + + return (final_result, mask_result) + + except Exception as e: + print(f"BlackBackgroundScanner error: {str(e)}") + empty_mask = torch.zeros_like(image) + return (image, empty_mask) \ No newline at end of file diff --git a/document_scanner.py b/document_scanner.py new file mode 100644 index 0000000..ab33dc0 --- /dev/null +++ b/document_scanner.py @@ -0,0 +1,200 @@ +import torch +import numpy as np +from .utils import ( + tensor_to_cv2, cv2_to_tensor, blank_page, to_grayscale, blur, + to_edges, find_vertices, crop_out, enhance_image +) + + +class DocumentScannerNode: + """ + ComfyUI node for document scanning with perspective correction and enhancement + """ + + CATEGORY = "image/processing" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "enhancement_method": (["sharpening", "cartooning", "clahe", "threshold", "adaptive_threshold", "flat_field"], { + "default": "sharpening" + }), + "edge_threshold_low": ("INT", { + "default": 20, + "min": 1, + "max": 100, + "step": 1 + }), + "edge_threshold_high": ("INT", { + "default": 70, + "min": 1, + "max": 255, + "step": 1 + }), + "blur_kernel_size": ("INT", { + "default": 5, + "min": 3, + "max": 15, + "step": 2 + }), + "skip_preprocessing": ("BOOLEAN", { + "default": False + }), + "return_debug_edges": ("BOOLEAN", { + "default": False + }) + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("scanned_image", "debug_edges") + FUNCTION = "scan_document" + + def scan_document(self, image, enhancement_method, edge_threshold_low, edge_threshold_high, + blur_kernel_size, skip_preprocessing, return_debug_edges): + """ + Main document scanning function + """ + try: + results = [] + debug_edges_batch = [] + + # Process each image in the batch + for i in range(image.shape[0]): + # Convert tensor to OpenCV format + cv2_image = tensor_to_cv2(image[i:i+1]) + + # Document scanning pipeline + processed_image, edges_debug = self._process_single_image( + cv2_image, enhancement_method, edge_threshold_low, + edge_threshold_high, blur_kernel_size, skip_preprocessing + ) + + # Convert back to tensor format + result_tensor = cv2_to_tensor(processed_image) + results.append(result_tensor) + + if return_debug_edges: + edges_tensor = cv2_to_tensor(edges_debug) + debug_edges_batch.append(edges_tensor) + + # Combine batch results + final_result = torch.cat(results, dim=0) + + if return_debug_edges and debug_edges_batch: + debug_result = torch.cat(debug_edges_batch, dim=0) + else: + # Return empty tensor with same batch size if no debug requested + debug_result = torch.zeros_like(final_result) + + return (final_result, debug_result) + + except Exception as e: + print(f"DocumentScanner error: {str(e)}") + # Fallback: return original image + empty_debug = torch.zeros_like(image) + return (image, empty_debug) + + def _process_single_image(self, cv2_image, enhancement_method, edge_threshold_low, + edge_threshold_high, blur_kernel_size, skip_preprocessing): + """ + Process a single image through the document scanning pipeline + """ + try: + original_image = cv2_image.copy() + + # Step 1: Preprocessing (optional) + if not skip_preprocessing: + processed_image = blank_page(cv2_image) + else: + processed_image = cv2_image + + # Step 2: Convert to grayscale + grayscale = to_grayscale(processed_image) + + # Step 3: Apply blur + blurred = blur(grayscale, blur_kernel_size) + + # Step 4: Edge detection + edges = to_edges(blurred, edge_threshold_low, edge_threshold_high) + + # Create debug visualization of edges + edges_debug = np.stack([edges, edges, edges], axis=2) # Convert to 3-channel for visualization + + # Step 5: Find document vertices + vertices = find_vertices(edges) + + # Step 6: Perspective correction + cropped = crop_out(original_image, vertices) + + # Step 7: Enhancement + enhanced = enhance_image(cropped, enhancement_method) + + return enhanced, edges_debug + + except Exception as e: + print(f"Error processing single image: {str(e)}") + # Fallback: return original image + edges_debug = np.zeros_like(cv2_image) + return cv2_image, edges_debug + + @classmethod + def IS_CHANGED(cls, **kwargs): + """ + This method can be used to determine if the node should be re-executed + For now, we'll always execute to ensure fresh results + """ + return float("NaN") + + +# Alternative node with simplified interface +class SimpleDocumentScannerNode: + """ + Simplified version of DocumentScanner with fewer parameters + """ + + CATEGORY = "image/processing" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "enhancement": (["auto", "text", "photo"], { + "default": "auto" + }) + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("scanned_image",) + FUNCTION = "simple_scan" + + def simple_scan(self, image, enhancement): + """ + Simplified document scanning with preset configurations + """ + # Map simple enhancement options to detailed methods + enhancement_map = { + "auto": "sharpening", + "text": "adaptive_threshold", + "photo": "clahe" + } + + enhancement_method = enhancement_map.get(enhancement, "sharpening") + + # Use DocumentScannerNode with default parameters + scanner = DocumentScannerNode() + result, _ = scanner.scan_document( + image=image, + enhancement_method=enhancement_method, + edge_threshold_low=20, + edge_threshold_high=70, + blur_kernel_size=5, + skip_preprocessing=False, + return_debug_edges=False + ) + + return (result,) \ No newline at end of file diff --git a/gitignore.txt b/gitignore.txt new file mode 100644 index 0000000..5156ea1 --- /dev/null +++ b/gitignore.txt @@ -0,0 +1,145 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# IDE files +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS generated files +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes +ehthumbs.db +Thumbs.db \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..0871fac --- /dev/null +++ b/utils.py @@ -0,0 +1,237 @@ +import cv2 +import numpy as np +import torch + + +def tensor_to_cv2(tensor_image): + """Convert ComfyUI tensor to OpenCV format""" + # ComfyUI images are (batch, height, width, channels) in RGB + # OpenCV expects (height, width, channels) in BGR + if len(tensor_image.shape) == 4: + # Take first image from batch + tensor_image = tensor_image[0] + + # Convert from tensor to numpy array (0-1 range to 0-255) + cv2_image = tensor_image.cpu().numpy() + cv2_image = (cv2_image * 255).astype(np.uint8) + + # Convert RGB to BGR for OpenCV + cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_RGB2BGR) + + return cv2_image + + +def cv2_to_tensor(cv2_image): + """Convert OpenCV image to ComfyUI tensor format""" + # Convert BGR to RGB + rgb_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) + + # Convert to float32 and normalize to 0-1 range + tensor_image = rgb_image.astype(np.float32) / 255.0 + + # Convert to torch tensor and add batch dimension + tensor_image = torch.from_numpy(tensor_image).unsqueeze(0) + + return tensor_image + + +def reorder(vertices): + """Reorder vertices to top-left, top-right, bottom-right, bottom-left""" + reordered = np.zeros_like(vertices, dtype=np.float32) + add = vertices.sum(1) + reordered[0] = vertices[np.argmin(add)] # top-left (smallest sum) + reordered[2] = vertices[np.argmax(add)] # bottom-right (largest sum) + diff = np.diff(vertices, axis=1) + reordered[1] = vertices[np.argmin(diff)] # top-right (smallest diff) + reordered[3] = vertices[np.argmax(diff)] # bottom-left (largest diff) + return reordered + + +def blank_page(im): + """Remove text using morphological operations and GrabCut""" + kernel = np.ones((5,5), np.uint8) + img = cv2.morphologyEx(im, cv2.MORPH_CLOSE, kernel, iterations=3) + + mask = np.zeros(img.shape[:2], np.uint8) + bgdModel = np.zeros((1,65), np.float64) + fgdModel = np.zeros((1,65), np.float64) + rect = (20, 20, img.shape[1]-20, img.shape[0]-20) + + try: + cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT) + mask2 = np.where((mask==2)|(mask==0), 0, 1).astype('uint8') + img = img * mask2[:,:,np.newaxis] + except: + # Fallback: return original if GrabCut fails + pass + + return img + + +def to_grayscale(im): + """Convert image to grayscale""" + return cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) + + +def blur(im, kernel_size=5): + """Apply bilateral filter for edge-preserving smoothing""" + return cv2.bilateralFilter(im, kernel_size, 60, 60) + + +def to_edges(im, low_threshold=20, high_threshold=70): + """Apply Canny edge detection""" + return cv2.Canny(im, low_threshold, high_threshold) + + +def find_vertices(im): + """Find document vertices using contour detection""" + contours, hierarchy = cv2.findContours(im, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + + h, w = im.shape + area = (w - 10) * (h - 10) + area_found = area * 0.5 + + # Default fallback vertices (full image) + contour = np.array([[1, 1], [1, h-1], [w-1, h-1], [w-1, 1]]) + + for i in contours: + perimeter = cv2.arcLength(i, True) + approx = cv2.approxPolyDP(i, 0.03 * perimeter, True) + + if (len(approx) == 4 and + cv2.isContourConvex(approx) and + area_found < cv2.contourArea(approx) < area): + area_found = cv2.contourArea(approx) + contour = approx + + return contour.reshape((4, 2)) + + +def crop_out(im, vertices): + """Apply perspective transform to crop document""" + vertices = reorder(vertices) + (a, b, c, d) = vertices + + # Calculate optimal output dimensions + w1 = np.sqrt(((c[0] - d[0]) ** 2) + ((c[1] - d[1]) ** 2)) + w2 = np.sqrt(((b[0] - a[0]) ** 2) + ((b[1] - a[1]) ** 2)) + width = max(int(w1), int(w2)) + + h1 = np.sqrt(((b[0] - c[0]) ** 2) + ((b[1] - c[1]) ** 2)) + h2 = np.sqrt(((a[0] - d[0]) ** 2) + ((a[1] - d[1]) ** 2)) + height = max(int(h1), int(h2)) + + # Define target rectangle + target = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]], + dtype="float32") + + try: + transform = cv2.getPerspectiveTransform(vertices, target) + return cv2.warpPerspective(im, transform, (width, height)) + except: + # Fallback: return original image if perspective transform fails + return im + + +def enhance_sharpening(im): + """Apply sharpening enhancement""" + kernel_sharpening = np.array([[0,-1,0], + [-1, 5,-1], + [0,-1,0]]) + sharpened = cv2.filter2D(im, -1, kernel_sharpening) + + value1, value2 = 30, 25 + hsv = cv2.cvtColor(sharpened, cv2.COLOR_BGR2HSV) + h, s, v = cv2.split(hsv) + + lim = 255 - value1 + v[v > lim] = 255 + v[v <= lim] += value1 + + lim = 255 - value2 + s[s > lim] = 255 + s[s <= lim] += value2 + + final_hsv = cv2.merge((h, s, v)) + return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR) + + +def enhance_cartooning(im): + """Apply cartooning enhancement""" + num_down = 2 + num_bilateral = 7 + + img_color = im + for _ in range(num_down): + img_color = cv2.pyrDown(img_color) + + for _ in range(num_bilateral): + img_color = cv2.bilateralFilter(img_color, d=9, sigmaColor=9, sigmaSpace=7) + + for _ in range(num_down): + img_color = cv2.pyrUp(img_color) + + img_color = img_color[:im.shape[0], :im.shape[1], :im.shape[2]] + + img_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) + img_blur = cv2.medianBlur(img_gray, 7) + + img_edge = cv2.adaptiveThreshold(img_blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C, + cv2.THRESH_BINARY, blockSize=15, C=3) + + img_edge = cv2.cvtColor(img_edge, cv2.COLOR_GRAY2BGR) + return cv2.bitwise_and(img_color, img_edge) + + +def enhance_clahe(im): + """Apply CLAHE enhancement""" + lab = cv2.cvtColor(im, cv2.COLOR_BGR2LAB) + l_channel, a, b = cv2.split(lab) + + clahe = cv2.createCLAHE(clipLimit=1, tileGridSize=(8,8)) + cl = clahe.apply(l_channel) + + limg = cv2.merge((cl, a, b)) + return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR) + + +def enhance_threshold(im): + """Apply thresholding with denoising""" + gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) + ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) + denoised = cv2.fastNlMeansDenoising(thresh, 11, 31, 9) + return cv2.cvtColor(denoised, cv2.COLOR_GRAY2BGR) + + +def enhance_adaptive_threshold(im): + """Apply adaptive thresholding with denoising""" + gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) + thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, + cv2.THRESH_BINARY, 23, 5) + denoised = cv2.fastNlMeansDenoising(thresh, 11, 31, 9) + return cv2.cvtColor(denoised, cv2.COLOR_GRAY2BGR) + + +def enhance_flat_field(im): + """Apply flat field correction""" + F = cv2.GaussianBlur(im, (401, 401), 0) + C = np.int64(np.round((im * np.mean(F)) / F)) + C = np.clip(C, 0, 255).astype(np.uint8) + return enhance_clahe(C) + + +def enhance_image(im, method="sharpening"): + """Apply enhancement based on selected method""" + enhancement_methods = { + "sharpening": enhance_sharpening, + "cartooning": enhance_cartooning, + "clahe": enhance_clahe, + "threshold": enhance_threshold, + "adaptive_threshold": enhance_adaptive_threshold, + "flat_field": enhance_flat_field + } + + if method in enhancement_methods: + return enhancement_methods[method](im) + else: + return enhance_sharpening(im) # Default fallback \ No newline at end of file