Update app.py

app.py

@@ -38,7 +38,7 @@ class Generator(nn.Module):
 
         self.exit = nn.Sequential(
             nn.Conv2d(64, 3, 3, 1, 1),
-            nn.Sigmoid()  # IMPORTANT → output should be [0,1]
+            nn.Sigmoid()  # output [0,1]
         )
 
     def forward(self, x):
@@ -58,7 +58,7 @@ model.load_state_dict(checkpoint['generator'])
 model.eval()
 
 # ---------------------------
-# TRANSFORM (SAFE VERSION)
+# TRANSFORM
 # ---------------------------
 
 transform = transforms.ToTensor()
@@ -79,41 +79,41 @@ def enhance_image(input_image):
     # Convert tensor → numpy
     output = output.squeeze().permute(1, 2, 0).cpu().numpy()
 
-    # ---------------------------
-    # AUTO RANGE FIX (CRITICAL)
-    # ---------------------------
-    if output.min() < 0:  
-        # Model output is [-1,1]
+    # Handle range
+    if output.min() < 0:
         output = (output + 1) / 2
 
-    # Clamp safely
     output = np.clip(output, 0, 1)
 
-    # Convert to image
     output_img = (output * 255).astype(np.uint8)
 
-    # Resize back
-    output_img = Image.fromarray(output_img).resize(original_size, Image.BICUBIC)
+    # Resize back to original size
+    output_img = Image.fromarray(output_img)
+    output_img = output_img.resize(original_size, Image.BICUBIC)
     output_img = np.array(output_img)
 
     # ---------------------------
-    # VERY LIGHT POST-PROCESSING
+    # FINAL SAFE POST-PROCESSING
     # ---------------------------
 
-    # Small denoise only
+    # Light smoothing
     output_img = cv2.GaussianBlur(output_img, (3, 3), 0)
 
-    # VERY LIGHT sharpen
+    # Proper sharpening
     kernel = np.array([
         [0, -1, 0],
-        [-1, 4, -1],
+        [-1, 5, -1],
         [0, -1, 0]
     ])
     output_img = cv2.filter2D(output_img, -1, kernel)
 
+    # Blend with original (VERY IMPORTANT)
+    original_np = np.array(img.resize(original_size))
+    output_img = cv2.addWeighted(original_np, 0.7, output_img, 0.3, 0)
+
     output_img = np.clip(output_img, 0, 255)
 
-    # Save file
+    # Save for download
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
     Image.fromarray(output_img).save(temp_file.name)