Spaces:
Running
Running
| import numpy as np | |
| from skimage import transform, util | |
| from sklearn.ensemble import RandomForestRegressor | |
| from skimage.util import view_as_windows | |
| from PIL import Image | |
| # CONFIGURATION | |
| PATCH_SIZE = (5, 5) | |
| STEP = 1 | |
| N_ESTIMATORS = 50 | |
| MAX_DEPTH = 20 | |
| SCALE_FACTOR = 2 | |
| def extract_patches(img, patch_size, step): | |
| patches = view_as_windows(img, patch_size, step) | |
| h, w = patches.shape[:2] | |
| return patches.reshape(h * w, -1) | |
| def train_rf(X, y): | |
| rf = RandomForestRegressor(n_estimators=N_ESTIMATORS, max_depth=MAX_DEPTH, n_jobs=-1) | |
| rf.fit(X, y) | |
| return rf | |
| def predict_and_reconstruct(model, lr_img, patch_size, step, out_shape): | |
| lr_patches = extract_patches(lr_img, patch_size, step) | |
| preds = model.predict(lr_patches) | |
| patch_h, patch_w = patch_size | |
| img_h = (lr_img.shape[0] - patch_h) // step + 1 | |
| img_w = (lr_img.shape[1] - patch_w) // step + 1 | |
| result = np.zeros(out_shape) | |
| weight = np.zeros(out_shape) | |
| idx = 0 | |
| for i in range(img_h): | |
| for j in range(img_w): | |
| patch = preds[idx].reshape(patch_h, patch_w) | |
| result[i*step:i*step+patch_h, j*step:j*step+patch_w] += patch | |
| weight[i*step:i*step+patch_h, j*step:j*step+patch_w] += 1 | |
| idx += 1 | |
| weight[weight == 0] = 1 | |
| return result / weight | |
| def random_forest_upscale(pil_img: Image.Image) -> Image.Image: | |
| img = np.array(pil_img) / 255.0 # Normalize | |
| if img.ndim == 2: | |
| img = np.expand_dims(img, axis=-1) | |
| hr_shape = (img.shape[0] * SCALE_FACTOR, img.shape[1] * SCALE_FACTOR) | |
| sr_channels = [] | |
| for c in range(img.shape[2]): | |
| channel = img[:, :, c] | |
| hr_channel = transform.resize(channel, hr_shape) | |
| lr_channel = transform.resize(hr_channel, (hr_shape[0] // SCALE_FACTOR, hr_shape[1] // SCALE_FACTOR)) | |
| lr_channel_up = transform.resize(lr_channel, hr_shape) | |
| X = extract_patches(lr_channel_up, PATCH_SIZE, STEP) | |
| y = extract_patches(hr_channel, PATCH_SIZE, STEP) | |
| rf_model = train_rf(X, y) | |
| sr = predict_and_reconstruct(rf_model, lr_channel_up, PATCH_SIZE, STEP, hr_shape) | |
| sr_channels.append(sr) | |
| sr_image = np.stack(sr_channels, axis=-1) | |
| sr_image = np.clip(sr_image * 255, 0, 255).astype(np.uint8) | |
| return Image.fromarray(sr_image) | |