CoolFace
Apppublic

ganeshkumar383/AI-Based-Image-Deblurring-App

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
color_preservation.py210 linesDownload Raw Back to modules
1"""
2Color Preservation Module
3========================
4
5Utilities to ensure perfect color preservation during image enhancement.
6Only sharpness, clarity, and focus should be improved while maintaining
7the exact original colors.
8"""
9
10import cv2
11import numpy as np
12from typing import Tuple, Optional
13import logging
14
15# Configure logging
16logging.basicConfig(level=logging.INFO)
17logger = logging.getLogger(__name__)
18
19class ColorPreserver:
20    """Utilities for preserving colors during image enhancement"""
21    
22    @staticmethod
23    def preserve_colors_during_enhancement(original: np.ndarray, 
24                                         enhanced: np.ndarray,
25                                         preservation_strength: float = 0.8) -> np.ndarray:
26        """
27        Preserve original colors while keeping enhancement benefits
28        
29        Args:
30            original: Original image (BGR)
31            enhanced: Enhanced image (BGR) 
32            preservation_strength: How much to preserve original colors (0-1)
33            
34        Returns:
35            np.ndarray: Color-preserved enhanced image
36        """
37        try:
38            # Convert to LAB color space for better color/brightness separation
39            original_lab = cv2.cvtColor(original, cv2.COLOR_BGR2LAB)
40            enhanced_lab = cv2.cvtColor(enhanced, cv2.COLOR_BGR2LAB)
41            
42            # Split LAB channels
43            orig_l, orig_a, orig_b = cv2.split(original_lab)
44            enh_l, enh_a, enh_b = cv2.split(enhanced_lab)
45            
46            # Keep enhanced brightness (L channel) but preserve original colors (A, B channels)
47            preserved_a = (preservation_strength * orig_a + 
48                          (1 - preservation_strength) * enh_a).astype(np.uint8)
49            preserved_b = (preservation_strength * orig_b + 
50                          (1 - preservation_strength) * enh_b).astype(np.uint8)
51            
52            # Combine preserved colors with enhanced brightness
53            result_lab = cv2.merge([enh_l, preserved_a, preserved_b])
54            
55            # Convert back to BGR
56            result = cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR)
57            
58            logger.info("Color preservation applied")
59            return result
60            
61        except Exception as e:
62            logger.error(f"Error in color preservation: {e}")
63            return enhanced
64    
65    @staticmethod
66    def enhance_sharpness_only(image: np.ndarray, 
67                              sharpening_strength: float = 0.3) -> np.ndarray:
68        """
69        Enhance only sharpness without affecting colors
70        
71        Args:
72            image: Input image (BGR)
73            sharpening_strength: Sharpening strength (0-1)
74            
75        Returns:
76            np.ndarray: Sharpness-enhanced image with preserved colors
77        """
78        try:
79            # Convert to float for precision
80            img_float = image.astype(np.float64)
81            
82            # Create a subtle sharpening kernel
83            kernel = np.array([[-0.05, -0.1, -0.05],
84                              [-0.1,  1.4,  -0.1],
85                              [-0.05, -0.1, -0.05]]) * sharpening_strength
86            
87            # Add identity for original preservation
88            kernel[1, 1] += (1 - sharpening_strength)
89            
90            # Apply sharpening filter
91            sharpened = cv2.filter2D(img_float, -1, kernel)
92            
93            # Ensure no clipping artifacts that change colors
94            result = np.clip(sharpened, 0, 255).astype(np.uint8)
95            
96            return result
97            
98        except Exception as e:
99            logger.error(f"Error in sharpness-only enhancement: {e}")
100            return image
101    
102    @staticmethod
103    def accurate_unsharp_masking(image: np.ndarray,
104                               sigma: float = 1.0,
105                               amount: float = 0.5) -> np.ndarray:
106        """
107        Apply unsharp masking with perfect color preservation
108        
109        Args:
110            image: Input image (BGR)
111            sigma: Gaussian blur sigma for mask
112            amount: Sharpening amount
113            
114        Returns:
115            np.ndarray: Sharpened image with preserved colors
116        """
117        try:
118            # Work in high precision
119            img_float = image.astype(np.float64)
120            
121            # Create Gaussian blur
122            blurred = cv2.GaussianBlur(img_float, (0, 0), sigma)
123            
124            # Create unsharp mask
125            mask = img_float - blurred
126            
127            # Apply mask with careful amount control
128            sharpened = img_float + amount * mask
129            
130            # Careful clipping to preserve color accuracy
131            result = np.clip(sharpened, 0, 255)
132            result = np.round(result).astype(np.uint8)
133            
134            return result
135            
136        except Exception as e:
137            logger.error(f"Error in accurate unsharp masking: {e}")
138            return image
139    
140    @staticmethod
141    def convert_for_display(image_bgr: np.ndarray) -> np.ndarray:
142        """
143        Convert BGR image to RGB for proper display in Streamlit
144        
145        Args:
146            image_bgr: Image in BGR format
147            
148        Returns:
149            np.ndarray: Image in RGB format for display
150        """
151        try:
152            if len(image_bgr.shape) == 3 and image_bgr.shape[2] == 3:
153                return cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
154            return image_bgr
155        except Exception as e:
156            logger.error(f"Error converting for display: {e}")
157            return image_bgr
158    
159    @staticmethod
160    def validate_color_preservation(original: np.ndarray, 
161                                  processed: np.ndarray,
162                                  tolerance: float = 5.0) -> dict:
163        """
164        Validate that colors are preserved during processing
165        
166        Args:
167            original: Original image
168            processed: Processed image  
169            tolerance: Acceptable color difference
170            
171        Returns:
172            dict: Validation results
173        """
174        try:
175            # Convert to LAB for perceptual color comparison
176            orig_lab = cv2.cvtColor(original, cv2.COLOR_BGR2LAB)
177            proc_lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
178            
179            # Calculate color differences (A and B channels only)
180            diff_a = np.mean(np.abs(orig_lab[:, :, 1].astype(np.float32) - 
181                                  proc_lab[:, :, 1].astype(np.float32)))
182            diff_b = np.mean(np.abs(orig_lab[:, :, 2].astype(np.float32) - 
183                                  proc_lab[:, :, 2].astype(np.float32)))
184            
185            color_diff = (diff_a + diff_b) / 2.0
186            
187            return {
188                'color_difference': float(color_diff),
189                'colors_preserved': color_diff <= tolerance,
190                'a_channel_diff': float(diff_a),
191                'b_channel_diff': float(diff_b),
192                'tolerance_used': tolerance
193            }
194            
195        except Exception as e:
196            logger.error(f"Error validating color preservation: {e}")
197            return {'colors_preserved': False, 'error': str(e)}
198
199# Convenience functions for easy use
200def preserve_colors(original: np.ndarray, enhanced: np.ndarray) -> np.ndarray:
201    """Preserve colors from original in enhanced image"""
202    return ColorPreserver.preserve_colors_during_enhancement(original, enhanced)
203
204def sharpen_only(image: np.ndarray, strength: float = 0.3) -> np.ndarray:
205    """Sharpen image without changing colors"""
206    return ColorPreserver.enhance_sharpness_only(image, strength)
207
208def display_convert(image_bgr: np.ndarray) -> np.ndarray:
209    """Convert BGR to RGB for display"""
210    return ColorPreserver.convert_for_display(image_bgr)