| """ | |
| Author: https://github.com/datnguyen-tien204 | |
| """ | |
| import numpy as np | |
| import torch | |
| import math | |
| PI = math.pi | |
| def symmetric_brightness_curve(hour): | |
| return 0.035 + (1.0 - 0.035) * (0.5 - 0.5 * np.cos(2 * PI * (hour / 24.0))) | |
| def rgb_curve_original(hour): | |
| t = (hour % 24.0) / 24.0 | |
| sunset = np.exp(-((t - 0.75) / 0.08) ** 2) | |
| sunrise = np.exp(-((t - 0.25) / 0.08) ** 2) | |
| night = np.exp(-((t - 0.0) / 0.10) ** 2) + np.exp(-((t - 1.0) / 0.10) ** 2) | |
| noon = np.exp(-((t - 0.5) / 0.12) ** 2) | |
| return [1.0 + 0.4 * sunset + 0.22 * sunrise - 0.35 * night, | |
| 1.0 - 0.1 * sunset - 0.06 * night, | |
| 1.0 - 0.5 * sunset + 0.3 * night + 0.15 * noon] | |
| def create_target_final_v2(day_img, style_img, depth_map, hour_val, adain_module, style_stats): | |
| device, dtype, B = day_img.device, day_img.dtype, day_img.shape[0] | |
| img_original = (day_img + 1.0) / 2.0 | |
| target_list = [] | |
| for i in range(B): | |
| h = hour_val[i].item() | |
| bright = symmetric_brightness_curve(h) | |
| rgb = rgb_curve_original(h) | |
| tint = torch.tensor(rgb, device=device, dtype=dtype).view(1, 3, 1, 1) | |
| tint = tint * (0.6 + 0.4 * style_stats["mean"][i].view(1, 3, 1, 1)) | |
| img_i = img_original[i:i + 1] | |
| target_base = img_i * tint | |
| if bright < 0.8: | |
| adain_out = (adain_module(day_img[i:i + 1], style_img[i:i + 1]) + 1.0) / 2.0 | |
| night_int = (1.0 - bright) ** 0.8 * 0.85 | |
| d_i = depth_map[i:i + 1] | |
| d_norm = (d_i - d_i.min()) / (d_i.max() - d_i.min() + 1e-6) | |
| sky_mask = torch.sigmoid((0.15 - d_norm) * 10.0) | |
| blend = night_int * (0.7 + 0.3 * sky_mask) | |
| target_mixed = target_base * (1 - blend) + adain_out * blend | |
| else: | |
| target_mixed = target_base | |
| sky_mask = torch.zeros_like(target_mixed[:, :1]) | |
| target_calc = target_mixed * bright | |
| if bright < 0.20: | |
| target_calc += img_i.mean(1, keepdim=True) * tint * 0.03 | |
| if bright < 0.8: target_calc *= (1.0 - sky_mask * 0.5) | |
| noon_w = torch.tensor(np.exp(-(abs(h - 12.0) ** 2) / 8.0), device=device, dtype=dtype) | |
| target_list.append(target_calc * (1 - noon_w) + img_i * noon_w) | |
| return (torch.cat(target_list, dim=0) * 2.0 - 1.0).clamp(-1, 1) |