Added EDI

app.py

@@ -18,56 +18,62 @@ from models.random_forest_sr import random_forest_upscale
 from PIL import Image
 import numpy as np
 
+from PIL import Image
+import numpy as np
+
 def edge_directed_interpolation(lr_img_pil, scale=2):
-    # Convert input PIL image to grayscale numpy array
-    # lr_img = np.array(lr_img_pil.convert("L"))
+    # Ensure input is a PIL Image
     if isinstance(lr_img_pil, np.ndarray):
         lr_img_pil = Image.fromarray(lr_img_pil)
 
-    lr_img = np.array(lr_img_pil.convert("L"))
-    h, w = lr_img.shape
-    hr_h, hr_w = h * scale, w * scale
+    # Convert to RGB
+    lr_img_rgb = lr_img_pil.convert("RGB")
+    lr_img = np.array(lr_img_rgb)
 
-    hr_img = np.zeros((hr_h, hr_w), dtype=np.uint8)
+    h, w, c = lr_img.shape
+    hr_h, hr_w = h * scale, w * scale
+    hr_img = np.zeros((hr_h, hr_w, c), dtype=np.uint8)
 
     # Copy original pixels to even positions
     for i in range(h):
         for j in range(w):
             hr_img[i * scale, j * scale] = lr_img[i, j]
 
-    # Interpolate diagonal pixels
+    # Interpolate diagonal pixels per channel
     for i in range(0, hr_h, scale):
         for j in range(0, hr_w, scale):
             if i + scale < hr_h and j + scale < hr_w:
-                p1 = hr_img[i, j]
-                p2 = hr_img[i, j + scale]
-                p3 = hr_img[i + scale, j]
-                p4 = hr_img[i + scale, j + scale]
+                for ch in range(c):
+                    p1 = hr_img[i, j, ch]
+                    p2 = hr_img[i, j + scale, ch]
+                    p3 = hr_img[i + scale, j, ch]
+                    p4 = hr_img[i + scale, j + scale, ch]
 
-                d1 = abs(int(p1) - int(p4))
-                d2 = abs(int(p2) - int(p3))
+                    d1 = abs(int(p1) - int(p4))
+                    d2 = abs(int(p2) - int(p3))
 
-                interp = (int(p1) + int(p4)) // 2 if d1 < d2 else (int(p2) + int(p3)) // 2
-                hr_img[i + scale // 2, j + scale // 2] = interp
+                    interp = (int(p1) + int(p4)) // 2 if d1 < d2 else (int(p2) + int(p3)) // 2
+                    hr_img[i + scale // 2, j + scale // 2, ch] = interp
 
-    # Fill remaining zero pixels
+    # Fill remaining zero pixels per channel
     for i in range(hr_h):
         for j in range(hr_w):
-            if hr_img[i, j] == 0:
-                neighbors = []
-                if i - 1 >= 0:
-                    neighbors.append(hr_img[i - 1, j])
-                if i + 1 < hr_h:
-                    neighbors.append(hr_img[i + 1, j])
-                if j - 1 >= 0:
-                    neighbors.append(hr_img[i, j - 1])
-                if j + 1 < hr_w:
-                    neighbors.append(hr_img[i, j + 1])
-                if neighbors:
-                    hr_img[i, j] = np.mean(neighbors).astype(np.uint8)
-
-    # Convert to PIL image (grayscale), then to RGB
-    return Image.fromarray(hr_img).convert("RGB")
+            for ch in range(c):
+                if hr_img[i, j, ch] == 0:
+                    neighbors = []
+                    if i - 1 >= 0:
+                        neighbors.append(hr_img[i - 1, j, ch])
+                    if i + 1 < hr_h:
+                        neighbors.append(hr_img[i + 1, j, ch])
+                    if j - 1 >= 0:
+                        neighbors.append(hr_img[i, j - 1, ch])
+                    if j + 1 < hr_w:
+                        neighbors.append(hr_img[i, j + 1, ch])
+                    if neighbors:
+                        hr_img[i, j, ch] = int(np.mean(neighbors))
+
+    return Image.fromarray(hr_img)
+
 
 
 # === Interfaces === #