Update app.py

fix color issues

app.py

@@ -1,4 +1,4 @@
-# app.py - Fixed Implementation with Color Correction
+# app.py - Complete Color-Corrected Implementation
 import os
 import cv2
 import time
@@ -10,67 +10,76 @@ import torch.nn.functional as F
 from PIL import Image
 from functools import partial
 
-# ====================== IMPROVED ARTIFACT MITIGATION ======================
-def balance_color_channels(image):
-    """Fix color channel imbalance causing yellow/green tint"""
-    b, g, r = cv2.split(image)
+# ====================== COLOR PRESERVATION FUNCTIONS ======================
+def preserve_original_colors(original, processed):
+    """Transfer colors from original to processed image"""
+    # Convert to LAB color space
+    original_lab = cv2.cvtColor(original, cv2.COLOR_RGB2LAB)
+    processed_lab = cv2.cvtColor(processed, cv2.COLOR_RGB2LAB)
     
-    # Calculate channel means to detect imbalance
-    b_mean, g_mean, r_mean = np.mean(b), np.mean(g), np.mean(r)
+    # Replace color information
+    processed_l, _, _ = cv2.split(processed_lab)
+    orig_l, orig_a, orig_b = cv2.split(original_lab)
     
-    # Detect yellow tint (low blue, high green/red)
-    if b_mean < (g_mean * 0.9) or b_mean < (r_mean * 0.9):
-        # Boost blue channel
-        b = np.clip(b * 1.2, 0, 255).astype(np.uint8)
-        # Reduce green slightly
-        g = np.clip(g * 0.9, 0, 255).astype(np.uint8)
+    # Use luminance from processed image but color from original
+    # Resize original color channels to match processed dimensions
+    h, w = processed_l.shape[:2]
+    resized_a = cv2.resize(orig_a, (w, h), interpolation=cv2.INTER_LINEAR)
+    resized_b = cv2.resize(orig_b, (w, h), interpolation=cv2.INTER_LINEAR)
     
-    return cv2.merge([b, g, r])
-
-def fix_chromatic_aberration(image):
-    """Improved color fringing reduction"""
-    # Reduce bilateral filter strength to preserve detail
-    filtered = cv2.bilateralFilter(image, d=3, sigmaColor=25, sigmaSpace=5)
-    # Fix color balance
-    return balance_color_channels(filtered)
+    # Create color corrected image
+    color_corrected = cv2.merge([processed_l, resized_a, resized_b])
+    return cv2.cvtColor(color_corrected, cv2.COLOR_LAB2RGB)
 
-def apply_anti_ringing(img):
-    """Reduce halo artifacts around edges"""
+def fix_color_cast(img):
+    """Remove color cast from image"""
+    # For grayscale/black&white images, force true grayscale
     gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    edges = cv2.Canny(gray, 100, 200)
-    dilated = cv2.dilate(edges, np.ones((2,2), np.uint8))  # Smaller kernel
+    gray_rgb = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
     
-    mask = cv2.GaussianBlur(dilated.astype(np.float32), (0,0), sigmaX=1.5)
-    mask = (mask / 255.0)[:,:,np.newaxis]
+    # Detect if image is likely grayscale
+    diff_r = np.abs(img[:,:,0].astype(np.float32) - gray.astype(np.float32))
+    diff_g = np.abs(img[:,:,1].astype(np.float32) - gray.astype(np.float32))
+    diff_b = np.abs(img[:,:,2].astype(np.float32) - gray.astype(np.float32))
     
-    filtered = cv2.bilateralFilter(img, d=2, sigmaColor=20, sigmaSpace=2)
-    return (img * (1-mask*0.7) + filtered * (mask*0.7)).astype(np.uint8)
-
-def color_normalize(image):
-    """Normalize colors to prevent shifts"""
-    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
-    l, a, b = cv2.split(lab)
+    total_diff = (np.mean(diff_r) + np.mean(diff_g) + np.mean(diff_b))/3
     
-    # Adjust a and b channels to prevent color shifts
-    a_mean, b_mean = np.mean(a), np.mean(b)
+    # If grayscale-like, force true grayscale
+    if total_diff < 10:  # Threshold for considering an image grayscale
+        return gray_rgb
     
-    # Pull toward neutral color if strong cast is detected
-    if abs(a_mean - 128) > 10 or abs(b_mean - 128) > 10:
-        a = np.clip(a * 0.85 + 128 * 0.15, 0, 255).astype(np.uint8)
-        b = np.clip(b * 0.85 + 128 * 0.15, 0, 255).astype(np.uint8)
+    # Otherwise just correct color balance
+    b, g, r = cv2.split(img)
+    r_avg, g_avg, b_avg = np.mean(r), np.mean(g), np.mean(b)
     
-    return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2RGB)
+    # Compute grayscale average
+    gray_avg = np.mean(gray)
+    
+    # Adjust channels to balance
+    r = np.clip(r * (gray_avg / r_avg) if r_avg > 0 else r, 0, 255).astype(np.uint8)
+    g = np.clip(g * (gray_avg / g_avg) if g_avg > 0 else g, 0, 255).astype(np.uint8)
+    b = np.clip(b * (gray_avg / b_avg) if b_avg > 0 else b, 0, 255).astype(np.uint8)
+    
+    return cv2.merge([b, g, r])
 
-def hybrid_upscale(image, neural_result, blend_factor=0.65):  # Reduced influence
-    """Blend neural and traditional upscaling with color preservation"""
-    traditional = cv2.resize(image, (neural_result.shape[1], neural_result.shape[0]), 
-                           interpolation=cv2.INTER_CUBIC)
+def simple_edge_enhance(img):
+    """Enhance edges without color distortion"""
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, 50, 150)
+    dilated = cv2.dilate(edges, np.ones((2,2), np.uint8))
+    
+    # Create edge mask
+    edge_mask = dilated.astype(np.float32) / 255.0
+    
+    # Sharpen image while preserving colors
+    blurred = cv2.GaussianBlur(img, (0, 0), 3)
+    sharpened = cv2.addWeighted(img, 1.5, blurred, -0.5, 0)
     
-    # Blend with reduced neural influence to preserve original colors
-    blended = cv2.addWeighted(neural_result, blend_factor, traditional, 1-blend_factor, 0)
+    # Apply sharpening only to edges
+    edge_mask = cv2.cvtColor(edge_mask[:,:,np.newaxis], cv2.COLOR_GRAY2RGB)
+    enhanced = img * (1 - edge_mask) + sharpened * edge_mask
     
-    # Normalize colors
-    return color_normalize(blended)
+    return enhanced.astype(np.uint8)
 
 # ====================== MODEL ARCHITECTURE ======================
 class SelfAttention(nn.Module):
@@ -117,6 +126,8 @@ class UltraEfficientSR(nn.Module):
         self.upconv2 = nn.Conv2d(64, 256, 3, padding=1)
         self.pixel_shuffle = nn.PixelShuffle(2)
         self.final = nn.Conv2d(64, 3, 3, padding=1)
+        
+        # Identity color preserving layer
         self.color_conv = nn.Conv2d(3, 3, 1)
         self._initialize_weights()
         
@@ -127,7 +138,7 @@ class UltraEfficientSR(nn.Module):
                 if m.bias is not None:
                     nn.init.zeros_(m.bias)
         
-        # Initialize color conv as identity for better color preservation
+        # Initialize color conv with identity matrix for color preservation
         with torch.no_grad():
             identity = torch.eye(3).reshape(3, 3, 1, 1)
             self.color_conv.weight.copy_(identity)
@@ -156,10 +167,20 @@ class UltraEfficientSR(nn.Module):
 
 # ====================== PROCESSING PIPELINE ======================
 def process_tile(model, tile, scale_factor):
+    # Preserve original for color reference
+    original_tile = tile.copy()
+    
+    # Process with model
     tile_tensor = torch.tensor(tile/255.0, dtype=torch.float32).permute(2,0,1).unsqueeze(0)
     with torch.no_grad():
         output = model(tile_tensor, scale_factor)
-    return output.squeeze().permute(1,2,0).clamp(0,1).numpy() * 255
+    
+    # Get raw output
+    raw_output = output.squeeze().permute(1,2,0).clamp(0,1).numpy() * 255
+    
+    # Color correction
+    color_corrected = preserve_original_colors(original_tile, raw_output.astype(np.uint8))
+    return color_corrected
 
 def create_pyramid_weights(h, w):
     y = np.linspace(0, 1, h)
@@ -186,7 +207,7 @@ def process_image_with_tiling(model, image, scale_factor, tile_size=256, overlap
             out_y1, out_x1 = y1*scale_factor, x1*scale_factor
             out_y2, out_x2 = y2*scale_factor, x2*scale_factor
             
-            # Fix weights size to match processed tile shape
+            # Create weights for this tile
             weights = create_pyramid_weights(processed.shape[0], processed.shape[1])
             
             output[out_y1:out_y2, out_x1:out_x2] += processed * weights
@@ -229,31 +250,39 @@ class CPUUpscaler:
             
         if image_np.shape[2] == 4:
             image_np = image_np[:,:,:3]
-            
+        
+        # Force grayscale for B&W images
+        is_grayscale = self._is_grayscale_image(image_np)
+        if is_grayscale:
+            gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
+            image_np = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
+        
         # Processing setup
         threads_used = self.energy_ctrl.adjust_processing(image_np.size)
         tile_size = self._calculate_optimal_tile_size(image_np)
         
+        # Save original for color reference
+        original_img = image_np.copy()
+        
         # Core processing
         if max(image_np.shape[:2]) > tile_size:
             output = process_image_with_tiling(self.model, image_np, scale_factor, tile_size)
         else:
             output = process_tile(self.model, image_np, scale_factor)
         
-        # Enhanced artifact mitigation pipeline
-        output = fix_chromatic_aberration(output)
-        output = apply_anti_ringing(output)
-        output = cv2.edgePreservingFilter(output, flags=cv2.NORMCONV_FILTER, 
-                                         sigma_s=40, sigma_r=0.3)
-        output = hybrid_upscale(image_np, output)
+        # Final color correction
+        output = preserve_original_colors(
+            cv2.resize(original_img, (output.shape[1], output.shape[0])), 
+            output
+        )
+        
+        # For B&W images, ensure true grayscale output
+        if is_grayscale:
+            gray = cv2.cvtColor(output, cv2.COLOR_RGB2GRAY)
+            output = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
         
-        # Final color correction for extreme cases
-        hsv = cv2.cvtColor(output, cv2.COLOR_RGB2HSV)
-        h, s, v = cv2.split(hsv)
-        # Reduce saturation in yellow regions
-        yellow_mask = np.logical_and(h > 20, h < 35)
-        s[yellow_mask] = s[yellow_mask] * 0.7
-        output = cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2RGB)
+        # Final edge enhancement
+        output = simple_edge_enhance(output)
         
         # Metrics
         metrics = {
@@ -261,10 +290,23 @@ class CPUUpscaler:
             "input_resolution": f"{image_np.shape[1]}x{image_np.shape[0]}",
             "output_resolution": f"{output.shape[1]}x{output.shape[0]}",
             "threads_used": threads_used,
-            "tile_size": tile_size
+            "tile_size": tile_size,
+            "color_preservation": "Active"
         }
         
         return Image.fromarray(output), metrics
+    
+    def _is_grayscale_image(self, img, threshold=5):
+        """Detect if an image is effectively grayscale"""
+        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+        b, g, r = cv2.split(img)
+        
+        diff_r = np.abs(r.astype(np.float32) - gray.astype(np.float32))
+        diff_g = np.abs(g.astype(np.float32) - gray.astype(np.float32))
+        diff_b = np.abs(b.astype(np.float32) - gray.astype(np.float32))
+        
+        total_diff = (np.mean(diff_r) + np.mean(diff_g) + np.mean(diff_b))/3
+        return total_diff < threshold
 
 # ====================== GRADIO INTERFACE ======================
 def create_interface():