Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py
+# app.py - Final Corrected Implementation
 import os
 import cv2
 import time
@@ -10,39 +10,31 @@ import torch.nn.functional as F
 from PIL import Image
 from functools import partial
 
-# --------------------------
-# Artifact Mitigation Functions
-# --------------------------
+# ====================== ARTIFACT MITIGATION FUNCTIONS ======================
 def fix_chromatic_aberration(image):
-    """Fix color fringing artifacts by aligning RGB channels"""
+    """Align RGB channels to reduce color fringing"""
     return cv2.bilateralFilter(image, d=5, sigmaColor=50, sigmaSpace=10)
 
 def apply_anti_ringing(img):
-    """Reduce ringing artifacts around high-contrast edges"""
+    """Reduce halo/ringing 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))
     
-    mask = dilated.astype(np.float32) / 255.0
-    mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=2)
-    mask = mask[:,:,np.newaxis]
+    mask = cv2.GaussianBlur(dilated.astype(np.float32), (0,0), sigmaX=2)
+    mask = (mask / 255.0)[:,:,np.newaxis]
     
     filtered = cv2.bilateralFilter(img, d=3, sigmaColor=25, sigmaSpace=3)
-    result = img * (1-mask) + filtered * mask
-    
-    return result.astype(np.uint8)
+    return (img * (1-mask) + filtered * mask).astype(np.uint8)
 
 def hybrid_upscale(image, neural_result, blend_factor=0.8):
     """Blend neural and traditional upscaling"""
     h, w = image.shape[:2]
-    target_h, target_w = neural_result.shape[:2]
-    
-    traditional = cv2.resize(image, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
+    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)
 
-# --------------------------
-# Model Components
-# --------------------------
+# ====================== MODEL ARCHITECTURE ======================
 class SelfAttention(nn.Module):
     def __init__(self, channels):
         super().__init__()
@@ -57,7 +49,7 @@ class SelfAttention(nn.Module):
         k = self.key(x).view(batch, c, -1).permute(0, 2, 1)
         v = self.value(x).view(batch, c, -1)
         
-        attention = F.softmax(torch.bmm(q.float(), k.float()) / (c ** 0.5), dim=2)
+        attention = F.softmax(torch.bmm(q, k) / (c**0.5), dim=2)
         out = torch.bmm(attention, v).view(batch, c, h, w)
         return self.gamma * out + x
 
@@ -70,24 +62,24 @@ class ResidualBlock(nn.Module):
         
     def forward(self, x):
         residual = x
-        out = self.relu(self.conv1(x))
-        out = self.conv2(out)
-        return self.relu(out + residual)
+        x = self.relu(self.conv1(x))
+        x = self.conv2(x)
+        return self.relu(x + residual)
 
 class UltraEfficientSR(nn.Module):
-    def __init__(self, scale_factor=2):
+    def __init__(self):
         super().__init__()
-        self.initial = nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.initial = nn.Conv2d(3, 64, 3, padding=1)
         self.blocks = nn.Sequential(
             ResidualBlock(64),
             SelfAttention(64),
-            ResidualBlock(64),
+            ResidualBlock(64)
         )
-        self.upconv1 = nn.Conv2d(64, 256, kernel_size=3, padding=1)
-        self.upconv2 = nn.Conv2d(64, 256, kernel_size=3, padding=1)
+        self.upconv1 = nn.Conv2d(64, 256, 3, padding=1)
+        self.upconv2 = nn.Conv2d(64, 256, 3, padding=1)
         self.pixel_shuffle = nn.PixelShuffle(2)
-        self.final = nn.Conv2d(64, 3, kernel_size=3, padding=1)
-        self.color_conv = nn.Conv2d(3, 3, kernel_size=1)
+        self.final = nn.Conv2d(64, 3, 3, padding=1)
+        self.color_conv = nn.Conv2d(3, 3, 1)
         self._initialize_weights()
         
     def _initialize_weights(self):
@@ -115,32 +107,27 @@ class UltraEfficientSR(nn.Module):
             x = self.pixel_shuffle(x)
         
         x = self.final(x)
-        x = self.color_conv(x)
-        return x
+        return self.color_conv(x)
 
-# --------------------------
-# Processing Pipeline
-# --------------------------
-def process_tile(model, tile, scale_factor=2):
-    tile_tensor = torch.tensor(tile/255.0, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
+# ====================== PROCESSING PIPELINE ======================
+def process_tile(model, tile, scale_factor):
+    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)
-    output = output.squeeze().permute(1, 2, 0).clamp(0, 1).numpy() * 255
-    return output.astype(np.uint8)
+    return output.squeeze().permute(1,2,0).clamp(0,1).numpy() * 255
 
 def create_pyramid_weights(h, w):
     y = np.linspace(0, 1, h)
     x = np.linspace(0, 1, w)
     xx, yy = np.meshgrid(x, y)
     weights = np.minimum(np.minimum(xx, 1-xx), np.minimum(yy, 1-yy))
-    return np.minimum(1.0, weights * 4)[:, :, np.newaxis]
+    return np.minimum(1.0, weights * 4)[:,:,np.newaxis]
 
-def process_image_with_tiling(model, image, scale_factor=2, tile_size=256, overlap=32):
+def process_image_with_tiling(model, image, scale_factor, tile_size=256, overlap=32):
     h, w, c = image.shape
-    tile_size = min(tile_size, h, w)
-    out_h, out_w = h * scale_factor, w * scale_factor
-    output = np.zeros((out_h, out_w, c), dtype=np.float32)
-    weight_map = np.zeros((out_h, out_w, c), dtype=np.float32)
+    out_h, out_w = h*scale_factor, w*scale_factor
+    output = np.zeros((out_h, out_w, c), np.float32)
+    weight_map = np.zeros_like(output)
     
     effective_step = tile_size - 2*overlap
     for y in range(0, h, effective_step):
@@ -151,71 +138,60 @@ def process_image_with_tiling(model, image, scale_factor=2, tile_size=256, overl
             tile = image[y1:y2, x1:x2]
             processed = process_tile(model, tile, scale_factor)
             
-            out_y1, out_x1 = y1 * scale_factor, x1 * scale_factor
-            out_y2, out_x2 = y2 * scale_factor, x2 * scale_factor
+            out_y1, out_x1 = y1*scale_factor, x1*scale_factor
+            out_y2, out_x2 = y2*scale_factor, x2*scale_factor
             
-            tile_weights = create_pyramid_weights(tile.shape[0] * scale_factor, 
-                                                tile.shape[1] * scale_factor)
+            weights = create_pyramid_weights(tile.shape[0]*scale_factor, 
+                                           tile.shape[1]*scale_factor)
             
-            output[out_y1:out_y2, out_x1:out_x2] += processed * tile_weights
-            weight_map[out_y1:out_y2, out_x1:out_x2] += tile_weights
+            output[out_y1:out_y2, out_x1:out_x2] += processed * weights
+            weight_map[out_y1:out_y2, out_x1:out_x2] += weights
     
     valid_mask = weight_map > 0
     output[valid_mask] /= weight_map[valid_mask]
     return output.astype(np.uint8)
 
-# --------------------------
-# Energy Management
-# --------------------------
+# ====================== CORE SYSTEM COMPONENTS ======================
 class EnergyController:
     def __init__(self):
         self.available_threads = os.cpu_count()
         
     def adjust_processing(self, image_size):
-        threads = max(1, min(self.available_threads, image_size // (1024**2) + 1))
+        threads = max(1, min(self.available_threads, image_size//(1024**2)+1))
         torch.set_num_threads(threads)
         return threads
 
-# --------------------------
-# Main Upscaler Class
-# --------------------------
 class CPUUpscaler:
     def __init__(self):
-        self.device = torch.device("cpu")
-        self.model = self._create_model()
+        self.model = torch.quantization.quantize_dynamic(
+            UltraEfficientSR(), {nn.Conv2d}, dtype=torch.qint8
+        ).eval()
         self.energy_ctrl = EnergyController()
         
-    def _create_model(self):
-        model = UltraEfficientSR()
-        model.eval()
-        return torch.quantization.quantize_dynamic(
-            model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8
-        )
-    
     def _calculate_optimal_tile_size(self, image):
         gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
         edge_density = cv2.Laplacian(gray, cv2.CV_64F).var()
-        
-        if edge_density > 500: return 128
-        elif edge_density > 200: return 256
-        else: return 384
+        return 128 if edge_density > 500 else 256 if edge_density > 200 else 384
     
     def upscale(self, image, scale_factor=2):
-        if image is None: return None, {"error": "No image provided"}
-        
         start_time = time.time()
-        image_np = np.array(image) if isinstance(image, Image.Image) else image
         
+        # Input handling
+        if isinstance(image, Image.Image):
+            image_np = np.array(image)
+        else:
+            image_np = image.copy()
+            
         if image_np.shape[2] == 4:
-            image_np = image_np[:, :, :3]
-        
+            image_np = image_np[:,:,:3]
+            
+        # Processing setup
         threads_used = self.energy_ctrl.adjust_processing(image_np.size)
         tile_size = self._calculate_optimal_tile_size(image_np)
         
+        # Core processing
         if max(image_np.shape[:2]) > tile_size:
-            output = process_image_with_tiling(
-                self.model, image_np, scale_factor, 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)
         
@@ -225,8 +201,9 @@ class CPUUpscaler:
         output = cv2.edgePreservingFilter(output, flags=cv2.NORMCONV_FILTER, sigma_s=60, sigma_r=0.4)
         output = hybrid_upscale(image_np, output)
         
+        # Metrics
         metrics = {
-            "processing_time": f"{time.time() - start_time:.2f}s",
+            "processing_time": f"{time.time()-start_time:.2f}s",
             "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,
@@ -235,46 +212,35 @@ class CPUUpscaler:
         
         return Image.fromarray(output), metrics
 
-# --------------------------
-# Gradio Interface
-# --------------------------
-CITATIONS = {
-    "main_model": {"title": "EfficientSR: Efficient Neural Super-Resolution...", "doi": "10.1109/CVPR52729.2024.00709"},
-    "sparse_attention": {"title": "SparseWin...", "doi": "10.1109/ICCV48922.2025.01207"},
-    "hybrid_quant": {"title": "Hybrid 4-8 Bit Quantization...", "doi": "10.1109/TPAMI.2025.3056721"}
-}
-
+# ====================== GRADIO INTERFACE ======================
 def create_interface():
     upscaler = CPUUpscaler()
     
     def process_image(input_img, scale_factor):
-        scale_map = {"2x": 2, "3x": 3, "4x": 4}
+        scale_map = {"2x":2, "3x":3, "4x":4}
         output_img, metrics = upscaler.upscale(input_img, scale_map[scale_factor])
         return output_img, [input_img, output_img], metrics
     
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# Advanced CPU-Optimized Image Upscaler")
+        gr.Markdown("# Professional Image Upscaler")
         with gr.Row():
             with gr.Column(scale=1):
-                input_img = gr.Image(label="Input Image", type="pil")
-                scale_factor = gr.Radio(["2x", "3x", "4x"], value="2x", label="Scale Factor")
+                input_img = gr.Image(label="Input", type="pil")
+                scale_factor = gr.Radio(["2x","3x","4x"], value="2x", label="Scale")
                 upscale_btn = gr.Button("Upscale", variant="primary")
                 
             with gr.Column(scale=2):
-                output_img = gr.Image(label="Upscaled Result", type="pil")
-                comparison = gr.Gallery(label="Before/After Comparison", columns=2, height="auto")
-                metrics = gr.JSON(label="Performance Metrics")
+                output_img = gr.Image(label="Result", type="pil")
+                comparison = gr.Gallery(columns=2, height="auto")
+                metrics = gr.JSON(label="Metrics")
         
         upscale_btn.click(
-            process_image, [input_img, scale_factor], [output_img, comparison, metrics]
+            process_image, 
+            [input_img, scale_factor], 
+            [output_img, comparison, metrics]
         )
         
-        with gr.Accordion("Technical Details", open=False):
-            gr.Markdown("## Implementation Details")
-            gr.JSON(CITATIONS, label="Academic References")
-    
     return demo
 
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+    create_interface().launch()