models/edge_directed_interpolation.py

import numpy as np
import cv2

def edge_directed_interpolation(lr_img_pil, scale=2):
    lr_img = np.array(lr_img_pil.convert("L"))
    h, w = lr_img.shape
    hr_h, hr_w = h * scale, w * scale

    hr_img = np.zeros((hr_h, hr_w), dtype=np.uint8)

    for i in range(h):
        for j in range(w):
            hr_img[i * scale, j * scale] = lr_img[i, j]

    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]

                d1 = abs(int(p1) - int(p4))
                d2 = abs(int(p2) - int(p3))

                if d1 < d2:
                    interp = (int(p1) + int(p4)) // 2
                else:
                    interp = (int(p2) + int(p3)) // 2

                hr_img[i + scale // 2, j + scale // 2] = interp

    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)

    hr_img_pil = Image.fromarray(hr_img)
    return hr_img_pil