Spaces:
Sleeping
Sleeping
| #%% | |
| import numpy as np | |
| from pathlib import Path | |
| from PIL import Image | |
| w0 = np.array( | |
| [ | |
| [ | |
| [ | |
| [-1.80224888e-02, 4.05534595e-01, 1.62569676e-02], | |
| [3.86995256e-01, -1.58828032e00, 6.36388600e-01], | |
| [-1.01250924e-01, 6.00708783e-01, -1.56698138e-01], | |
| ], | |
| [ | |
| [-2.71480846e00, -1.42060733e00, -2.68766713e00], | |
| [-1.52638519e00, 1.73007298e01, -1.35628653e00], | |
| [-2.54135966e00, -1.04333639e00, -2.83616400e00], | |
| ], | |
| [ | |
| [1.95330644e00, 2.82546377e00, 1.74052000e00], | |
| [2.87443542e00, -1.98198109e01, 3.01367259e00], | |
| [1.84999144e00, 2.56530333e00, 1.73608303e00], | |
| ], | |
| ], | |
| [ | |
| [ | |
| [-2.12553531e-01, -1.84073091e00, -3.59450318e-02], | |
| [-1.36595523e00, 6.61149931e00, -1.21200264e00], | |
| [-1.61612287e-01, -1.88850796e00, -4.96574976e-02], | |
| ], | |
| [ | |
| [-1.69264674e00, 7.85137057e-01, -1.77036440e00], | |
| [3.79421353e-01, 5.04367065e00, 3.15313727e-01], | |
| [-1.68788207e00, 9.53136027e-01, -1.84203279e00], | |
| ], | |
| [ | |
| [1.44280708e00, 2.15313840e00, 1.39907610e00], | |
| [2.01263666e00, -1.40174408e01, 1.93493605e00], | |
| [1.39259851e00, 2.11386871e00, 1.38634503e00], | |
| ], | |
| ], | |
| ], | |
| ) | |
| b0 = np.array([-0.3287916, 0.75910544]) | |
| w1 = np.array([[[[53.911407]], [[-38.45282]]]]) | |
| b1 =np.array([15.667337]) | |
| threshold = 0.3323521614074707 | |
| def gelu(x): # tanh approximation — ~1e-3 off from torch's exact erf-based GELU | |
| return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3))) | |
| def conv3x3_reflect(x : np.ndarray, w: np.ndarray, b: np.ndarray): | |
| """x:(Cin,H,W) w:(Cout,Cin,3,3) b:(Cout,) -> (Cout,H,W)""" | |
| Cin, H, W = x.shape | |
| Cout = w.shape[0] | |
| xp = np.pad(x, ((0, 0), (1, 1), (1, 1)), mode="reflect") | |
| win : np.ndarray = np.lib.stride_tricks.sliding_window_view(xp, (3, 3), axis=(-2, -1)) | |
| win = win.transpose(1, 2, 0, 3, 4).reshape(H * W, Cin * 9) | |
| out = win @ w.reshape(Cout, -1).T + b | |
| return out.reshape(H, W, Cout).transpose(2, 0, 1) | |
| def conv1x1(x: np.ndarray, w: np.ndarray, b: np.ndarray): | |
| """x:(Cin,H,W) w:(Cout,Cin,1,1) b:(Cout,) -> (Cout,H,W)""" | |
| Cin, H, W = x.shape | |
| Cout = w.shape[0] | |
| return (w.reshape(Cout, Cin) @ x.reshape(Cin, -1) + b[:, None]).reshape(Cout, H, W) | |
| def predict(input_image): | |
| """img_rgb_u8: (H, W, 3) uint8 -> (H, W) logits.""" | |
| img_rgb_u8 = np.asarray(input_image.convert("RGB")) | |
| x = img_rgb_u8.astype(np.float32).transpose(2, 0, 1) / 255.0 | |
| h = conv3x3_reflect(x, w0, b0) | |
| h = gelu(h) | |
| logits= conv1x1(h, w1, b1)[0] | |
| pred = logits > threshold | |
| return pred | |
| def fix(input_image): | |
| pred = predict(input_image) | |
| x = np.asarray(input_image.convert("RGB")) | |
| # Replace predicted-positive pixels with the average of their non-predicted | |
| # neighbours in a 3x3 window (zero-padded — matches the torch version). | |
| keep = (~pred).astype(np.float32) # (H, W) | |
| xn = x.astype(np.float32) / 255.0 # (H, W, 3) in [0, 1] | |
| x_t = (xn * keep[..., None]).transpose(2, 0, 1) # (3, H, W) | |
| x_p = np.pad(x_t, ((0, 0), (1, 1), (1, 1))) # zero pad | |
| k_p = np.pad(keep, ((1, 1), (1, 1))) | |
| wx :np.ndarray= np.lib.stride_tricks.sliding_window_view(x_p, (3, 3), axis=(-2, -1)) # (3,H,W,3,3) | |
| wk = np.lib.stride_tricks.sliding_window_view(k_p, (3, 3)) # (H,W,3,3) | |
| sums = wx.sum(axis=(-2, -1)).transpose(1, 2, 0) # (H, W, 3) | |
| counts = wk.sum(axis=(-2, -1)) # (H, W) | |
| avg = sums / np.maximum(counts, 1)[..., None] | |
| out = np.where(pred[..., None], avg, xn) | |
| img = Image.fromarray(np.clip(out * 255, 0, 255).astype(np.uint8)) | |
| return img | |
| if __name__ == '__main__': | |
| path = Path("val4.png") | |
| input_image = Image.open(path) | |
| img = fix(input_image) | |
| img.save('NP.png') | |
| # %% | |