Update app.py

yellow tint fix

app.py

@@ -1,4 +1,4 @@
-# app.py - Final Corrected Implementation
+# app.py - Fixed Implementation with Color Correction
 import os
 import cv2
 import time
@@ -10,29 +10,67 @@ import torch.nn.functional as F
 from PIL import Image
 from functools import partial
 
-# ====================== ARTIFACT MITIGATION FUNCTIONS ======================
+# ====================== IMPROVED ARTIFACT MITIGATION ======================
+def balance_color_channels(image):
+    """Fix color channel imbalance causing yellow/green tint"""
+    b, g, r = cv2.split(image)
+    
+    # Calculate channel means to detect imbalance
+    b_mean, g_mean, r_mean = np.mean(b), np.mean(g), np.mean(r)
+    
+    # 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)
+    
+    return cv2.merge([b, g, r])
+
 def fix_chromatic_aberration(image):
-    """Align RGB channels to reduce color fringing"""
-    return cv2.bilateralFilter(image, d=5, sigmaColor=50, sigmaSpace=10)
+    """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)
 
 def apply_anti_ringing(img):
-    """Reduce halo/ringing artifacts around edges"""
+    """Reduce halo artifacts around edges"""
     gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
     edges = cv2.Canny(gray, 100, 200)
-    dilated = cv2.dilate(edges, np.ones((3,3), np.uint8))
+    dilated = cv2.dilate(edges, np.ones((2,2), np.uint8))  # Smaller kernel
     
-    mask = cv2.GaussianBlur(dilated.astype(np.float32), (0,0), sigmaX=2)
+    mask = cv2.GaussianBlur(dilated.astype(np.float32), (0,0), sigmaX=1.5)
     mask = (mask / 255.0)[:,:,np.newaxis]
     
-    filtered = cv2.bilateralFilter(img, d=3, sigmaColor=25, sigmaSpace=3)
-    return (img * (1-mask) + filtered * mask).astype(np.uint8)
+    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)
+    
+    # Adjust a and b channels to prevent color shifts
+    a_mean, b_mean = np.mean(a), np.mean(b)
+    
+    # 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)
+    
+    return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2RGB)
 
-def hybrid_upscale(image, neural_result, blend_factor=0.8):
-    """Blend neural and traditional upscaling"""
-    h, w = image.shape[:2]
+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)
-    return cv2.addWeighted(neural_result, blend_factor, traditional, 1-blend_factor, 0)
+    
+    # Blend with reduced neural influence to preserve original colors
+    blended = cv2.addWeighted(neural_result, blend_factor, traditional, 1-blend_factor, 0)
+    
+    # Normalize colors
+    return color_normalize(blended)
 
 # ====================== MODEL ARCHITECTURE ======================
 class SelfAttention(nn.Module):
@@ -88,6 +126,13 @@ class UltraEfficientSR(nn.Module):
                 nn.init.kaiming_normal_(m.weight, mode='fan_out')
                 if m.bias is not None:
                     nn.init.zeros_(m.bias)
+        
+        # Initialize color conv as identity for better color preservation
+        with torch.no_grad():
+            identity = torch.eye(3).reshape(3, 3, 1, 1)
+            self.color_conv.weight.copy_(identity)
+            if self.color_conv.bias is not None:
+                self.color_conv.bias.zero_()
     
     def forward(self, x, scale_factor=2):
         x = self.initial(x)
@@ -141,8 +186,8 @@ 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
             
-            weights = create_pyramid_weights(tile.shape[0]*scale_factor, 
-                                           tile.shape[1]*scale_factor)
+            # Fix weights size to match processed tile shape
+            weights = create_pyramid_weights(processed.shape[0], processed.shape[1])
             
             output[out_y1:out_y2, out_x1:out_x2] += processed * weights
             weight_map[out_y1:out_y2, out_x1:out_x2] += weights
@@ -195,12 +240,21 @@ class CPUUpscaler:
         else:
             output = process_tile(self.model, image_np, scale_factor)
         
-        # Artifact mitigation pipeline
+        # Enhanced artifact mitigation pipeline
         output = fix_chromatic_aberration(output)
         output = apply_anti_ringing(output)
-        output = cv2.edgePreservingFilter(output, flags=cv2.NORMCONV_FILTER, sigma_s=60, sigma_r=0.4)
+        output = cv2.edgePreservingFilter(output, flags=cv2.NORMCONV_FILTER, 
+                                         sigma_s=40, sigma_r=0.3)
         output = hybrid_upscale(image_np, output)
         
+        # 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)
+        
         # Metrics
         metrics = {
             "processing_time": f"{time.time()-start_time:.2f}s",