Add files via upload
This commit is contained in:
HarishCM
@@ -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
|
||||
+16
@@ -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']
|
||||
@@ -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)
|
||||
@@ -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,)
|
||||
+145
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user