Update app.py

app.py

@@ -58,16 +58,15 @@ model.load_state_dict(checkpoint['generator'])
 model.eval()
 
 # ---------------------------
-# TRANSFORM
+# TRANSFORM (FIXED)
 # ---------------------------
 
 transform = transforms.Compose([
-    transforms.Resize((128, 128)),
     transforms.ToTensor()
 ])
 
 # ---------------------------
-# INFERENCE FUNCTION
+# INFERENCE FUNCTION (FIXED)
 # ---------------------------
 
 def enhance_image(input_image):
@@ -79,38 +78,50 @@ def enhance_image(input_image):
     with torch.no_grad():
         output = model(input_tensor)
 
-    # Model output
-    output_img = output.squeeze().permute(1,2,0).cpu().numpy()
+    # Convert output tensor → image
+    output_img = output.squeeze().permute(1, 2, 0).cpu().numpy()
     output_img = (output_img * 255).astype(np.uint8)
 
-    # Resize back
-    output_img = Image.fromarray(output_img)
-    output_img = output_img.resize(original_size, Image.BICUBIC)
+    # Resize to original size
+    output_img = Image.fromarray(output_img).resize(original_size, Image.BICUBIC)
     output_img = np.array(output_img)
 
-    # 🔥 ORIGINAL IMAGE (for blending)
+    # Original image
     original_np = np.array(img)
 
-    # 🔥 STEP 1: LIGHT SMOOTHING (remove artifacts)
-    output_img = cv2.GaussianBlur(output_img, (3,3), 0)
+    # ---------------------------
+    # FIXED POST-PROCESSING
+    # ---------------------------
 
-    # 🔥 STEP 2: BLEND WITH ORIGINAL (VERY IMPORTANT)
-    final_img = cv2.addWeighted(original_np, 0.6, output_img, 0.4, 0)
+    # Convert RGB → BGR (IMPORTANT FIX)
+    original_bgr = cv2.cvtColor(original_np, cv2.COLOR_RGB2BGR)
+    output_bgr = cv2.cvtColor(output_img, cv2.COLOR_RGB2BGR)
 
-    # 🔥 STEP 3: VERY LIGHT SHARPEN (safe)
-    sharpen_kernel = np.array([[0,-1,0],
-                               [-1,5,-1],
-                               [0,-1,0]])
+    # Light smoothing
+    output_bgr = cv2.GaussianBlur(output_bgr, (3, 3), 0)
 
-    final_img = cv2.filter2D(final_img, -1, sharpen_kernel)
+    # Blend (more natural)
+    blended = cv2.addWeighted(original_bgr, 0.7, output_bgr, 0.3, 0)
 
-    final_img = np.clip(final_img, 0, 255)
+    # Mild sharpening (safe)
+    kernel = np.array([
+        [0, -1, 0],
+        [-1, 4, -1],
+        [0, -1, 0]
+    ])
+    sharpened = cv2.filter2D(blended, -1, kernel)
+
+    # Convert back BGR → RGB
+    final_img = cv2.cvtColor(sharpened, cv2.COLOR_BGR2RGB)
+
+    final_img = np.clip(final_img, 0, 255).astype(np.uint8)
 
     # Save for download
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
     Image.fromarray(final_img).save(temp_file.name)
 
     return final_img, temp_file.name
+
 # ---------------------------
 # GRADIO UI
 # ---------------------------