| """
|
| E3Diff: High-Resolution SAR-to-Optical Translation
|
| HuggingFace Spaces Deployment
|
|
|
| Features:
|
| - Full resolution processing with seamless tiling
|
| - Proper diffusion sampling (matching local inference)
|
| - TIFF output support
|
| """
|
|
|
| import os
|
| import sys
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import numpy as np
|
| from PIL import Image, ImageEnhance
|
| import gradio as gr
|
| from pathlib import Path
|
| import tempfile
|
| import time
|
| from functools import partial
|
| from huggingface_hub import hf_hub_download
|
|
|
|
|
| try:
|
| import spaces
|
| GPU_AVAILABLE = True
|
| except ImportError:
|
| GPU_AVAILABLE = False
|
| spaces = None
|
|
|
|
|
|
|
|
|
|
|
| def soft_pool2d(x, kernel_size=(2, 2), stride=None, force_inplace=False):
|
| if stride is None:
|
| stride = kernel_size
|
| if isinstance(kernel_size, int):
|
| kernel_size = (kernel_size, kernel_size)
|
| if isinstance(stride, int):
|
| stride = (stride, stride)
|
|
|
| batch, channels, height, width = x.shape
|
| kh, kw = kernel_size
|
| sh, sw = stride
|
| out_h = (height - kh) // sh + 1
|
| out_w = (width - kw) // sw + 1
|
|
|
| x_unfold = F.unfold(x, kernel_size=kernel_size, stride=stride)
|
| x_unfold = x_unfold.view(batch, channels, kh * kw, out_h * out_w)
|
| x_max = x_unfold.max(dim=2, keepdim=True)[0]
|
| exp_x = torch.exp(x_unfold - x_max)
|
| softpool = (x_unfold * exp_x).sum(dim=2) / (exp_x.sum(dim=2) + 1e-8)
|
| return softpool.view(batch, channels, out_h, out_w)
|
|
|
|
|
| class SoftPool2d(nn.Module):
|
| def __init__(self, kernel_size=(2, 2), stride=None, force_inplace=False):
|
| super(SoftPool2d, self).__init__()
|
| self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
|
| self.stride = stride if stride is not None else self.kernel_size
|
|
|
| def forward(self, x):
|
| return soft_pool2d(x, self.kernel_size, self.stride)
|
|
|
|
|
|
|
| class SoftPoolModule:
|
| soft_pool2d = staticmethod(soft_pool2d)
|
| SoftPool2d = SoftPool2d
|
| sys.modules['SoftPool'] = SoftPoolModule()
|
|
|
|
|
|
|
|
|
|
|
| import math
|
| from inspect import isfunction
|
|
|
| def exists(x):
|
| return x is not None
|
|
|
| def default(val, d):
|
| if exists(val):
|
| return val
|
| return d() if isfunction(d) else d
|
|
|
|
|
| class PositionalEncoding(nn.Module):
|
| def __init__(self, dim):
|
| super().__init__()
|
| self.dim = dim
|
|
|
| def forward(self, noise_level):
|
| count = self.dim // 2
|
| step = torch.arange(count, dtype=noise_level.dtype, device=noise_level.device) / count
|
| encoding = noise_level.unsqueeze(1) * torch.exp(-math.log(1e4) * step.unsqueeze(0))
|
| encoding = torch.cat([torch.sin(encoding), torch.cos(encoding)], dim=-1)
|
| return encoding
|
|
|
|
|
| class Swish(nn.Module):
|
| def forward(self, x):
|
| return x * torch.sigmoid(x)
|
|
|
|
|
| class FeatureWiseAffine(nn.Module):
|
| def __init__(self, in_channels, out_channels, use_affine_level=False):
|
| super(FeatureWiseAffine, self).__init__()
|
| self.use_affine_level = use_affine_level
|
| self.noise_func = nn.Sequential(nn.Linear(in_channels, out_channels*(1+self.use_affine_level)))
|
|
|
| def forward(self, x, noise_embed):
|
| batch = x.shape[0]
|
| if self.use_affine_level:
|
| gamma, beta = self.noise_func(noise_embed).view(batch, -1, 1, 1).chunk(2, dim=1)
|
| x = (1 + gamma) * x + beta
|
| else:
|
| x = x + self.noise_func(noise_embed).view(batch, -1, 1, 1)
|
| return x
|
|
|
|
|
| class Upsample(nn.Module):
|
| def __init__(self, dim):
|
| super().__init__()
|
| self.up = nn.Upsample(scale_factor=2, mode="nearest")
|
| self.conv = nn.Conv2d(dim, dim, 3, padding=1)
|
|
|
| def forward(self, x):
|
| return self.conv(self.up(x))
|
|
|
|
|
| class Downsample(nn.Module):
|
| def __init__(self, dim):
|
| super().__init__()
|
| self.conv = nn.Conv2d(dim, dim, 3, 2, 1)
|
|
|
| def forward(self, x):
|
| return self.conv(x)
|
|
|
|
|
| class Block(nn.Module):
|
| def __init__(self, dim, dim_out, groups=32, dropout=0, stride=1):
|
| super().__init__()
|
| self.block = nn.Sequential(
|
| nn.GroupNorm(groups, dim),
|
| Swish(),
|
| nn.Dropout(dropout) if dropout != 0 else nn.Identity(),
|
| nn.Conv2d(dim, dim_out, 3, stride=stride, padding=1)
|
| )
|
|
|
| def forward(self, x):
|
| return self.block(x)
|
|
|
|
|
| class ResnetBlock(nn.Module):
|
| def __init__(self, dim, dim_out, noise_level_emb_dim=None, dropout=0, use_affine_level=False, norm_groups=32):
|
| super().__init__()
|
| self.noise_func = FeatureWiseAffine(noise_level_emb_dim, dim_out, use_affine_level)
|
| self.c_func = nn.Conv2d(dim_out, dim_out, 1)
|
| self.block1 = Block(dim, dim_out, groups=norm_groups)
|
| self.block2 = Block(dim_out, dim_out, groups=norm_groups, dropout=dropout)
|
| self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
|
|
|
| def forward(self, x, time_emb, c):
|
| h = self.block1(x)
|
| h = self.noise_func(h, time_emb)
|
| h = self.block2(h)
|
|
|
| if c.shape[2:] != h.shape[2:]:
|
| c = F.interpolate(c, size=h.shape[2:], mode='bilinear', align_corners=False)
|
| h = self.c_func(c) + h
|
| return h + self.res_conv(x)
|
|
|
|
|
| class SelfAttention(nn.Module):
|
| def __init__(self, in_channel, n_head=1, norm_groups=32):
|
| super().__init__()
|
| self.n_head = n_head
|
| self.norm = nn.GroupNorm(norm_groups, in_channel)
|
| self.qkv = nn.Conv2d(in_channel, in_channel * 3, 1, bias=False)
|
| self.out = nn.Conv2d(in_channel, in_channel, 1)
|
|
|
| def forward(self, input, t=None, save_flag=None, file_num=None):
|
| batch, channel, height, width = input.shape
|
| n_head = self.n_head
|
| head_dim = channel // n_head
|
| norm = self.norm(input)
|
| qkv = self.qkv(norm).view(batch, n_head, head_dim * 3, height, width)
|
| query, key, value = qkv.chunk(3, dim=2)
|
| attn = torch.einsum("bnchw, bncyx -> bnhwyx", query, key).contiguous() / math.sqrt(channel)
|
| attn = attn.view(batch, n_head, height, width, -1)
|
| attn = torch.softmax(attn, -1)
|
| attn = attn.view(batch, n_head, height, width, height, width)
|
| out = torch.einsum("bnhwyx, bncyx -> bnchw", attn, value).contiguous()
|
| out = self.out(out.view(batch, channel, height, width))
|
| return out + input
|
|
|
|
|
| class ResnetBlocWithAttn(nn.Module):
|
| def __init__(self, dim, dim_out, *, noise_level_emb_dim=None, norm_groups=32, dropout=0, with_attn=False, size=256):
|
| super().__init__()
|
| self.with_attn = with_attn
|
| self.res_block = ResnetBlock(dim, dim_out, noise_level_emb_dim, norm_groups=norm_groups, dropout=dropout)
|
| if with_attn:
|
| self.attn = SelfAttention(dim_out, norm_groups=norm_groups)
|
|
|
| def forward(self, x, time_emb, c):
|
| x = self.res_block(x, time_emb, c)
|
| if self.with_attn:
|
| x = self.attn(x, time_emb)
|
| return x
|
|
|
|
|
|
|
| class CPEN(nn.Module):
|
| def __init__(self, inchannel=3):
|
| super(CPEN, self).__init__()
|
| from SoftPool import SoftPool2d
|
|
|
| self.conv1 = nn.Conv2d(inchannel, 64, 3, 1, 1)
|
| self.pool1 = SoftPool2d(kernel_size=(2, 2), stride=(2, 2))
|
| self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
|
| self.pool2 = SoftPool2d(kernel_size=(2, 2), stride=(2, 2))
|
| self.conv3 = nn.Conv2d(128, 256, 3, 1, 1)
|
| self.pool3 = SoftPool2d(kernel_size=(2, 2), stride=(2, 2))
|
| self.conv4 = nn.Conv2d(256, 512, 3, 1, 1)
|
| self.pool4 = SoftPool2d(kernel_size=(2, 2), stride=(2, 2))
|
| self.conv5 = nn.Conv2d(512, 1024, 3, 1, 1)
|
|
|
| def forward(self, x):
|
| c1 = self.pool1(F.leaky_relu(self.conv1(x)))
|
| c2 = self.pool2(F.leaky_relu(self.conv2(c1)))
|
| c3 = self.pool3(F.leaky_relu(self.conv3(c2)))
|
| c4 = self.pool4(F.leaky_relu(self.conv4(c3)))
|
| c5 = F.leaky_relu(self.conv5(c4))
|
| return c1, c2, c3, c4, c5
|
|
|
|
|
| class UNet(nn.Module):
|
| def __init__(self, in_channel=6, out_channel=3, inner_channel=32, norm_groups=32,
|
| channel_mults=(1, 2, 4, 8, 8), attn_res=(8,), res_blocks=3, dropout=0,
|
| with_noise_level_emb=True, image_size=128, condition_ch=3):
|
| super().__init__()
|
|
|
| self.res_blocks = res_blocks
|
| noise_level_channel = inner_channel
|
| self.noise_level_mlp = nn.Sequential(
|
| PositionalEncoding(inner_channel),
|
| nn.Linear(inner_channel, inner_channel * 4),
|
| Swish(),
|
| nn.Linear(inner_channel * 4, inner_channel)
|
| ) if with_noise_level_emb else None
|
|
|
| num_mults = len(channel_mults)
|
| pre_channel = inner_channel
|
| feat_channels = [pre_channel]
|
| now_res = image_size
|
|
|
| downs = [nn.Conv2d(in_channel, inner_channel, kernel_size=3, padding=1)]
|
| for ind in range(num_mults):
|
| is_last = (ind == num_mults - 1)
|
| use_attn = (now_res in attn_res)
|
| channel_mult = inner_channel * channel_mults[ind]
|
| for _ in range(0, res_blocks):
|
| downs.append(ResnetBlocWithAttn(pre_channel, channel_mult, noise_level_emb_dim=noise_level_channel,
|
| norm_groups=norm_groups, dropout=dropout, with_attn=use_attn, size=now_res))
|
| feat_channels.append(channel_mult)
|
| pre_channel = channel_mult
|
| if not is_last:
|
| downs.append(Downsample(pre_channel))
|
| feat_channels.append(pre_channel)
|
| now_res = now_res // 2
|
| self.downs = nn.ModuleList(downs)
|
|
|
| self.mid = nn.ModuleList([
|
| ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel,
|
| norm_groups=norm_groups, dropout=dropout, with_attn=True, size=now_res),
|
| ResnetBlocWithAttn(pre_channel, pre_channel, noise_level_emb_dim=noise_level_channel,
|
| norm_groups=norm_groups, dropout=dropout, with_attn=False, size=now_res)
|
| ])
|
|
|
| ups = []
|
| for ind in reversed(range(num_mults)):
|
| is_last = (ind < 1)
|
| use_attn = (now_res in attn_res)
|
| channel_mult = inner_channel * channel_mults[ind]
|
| for _ in range(0, res_blocks + 1):
|
| ups.append(ResnetBlocWithAttn(pre_channel + feat_channels.pop(), channel_mult,
|
| noise_level_emb_dim=noise_level_channel, norm_groups=norm_groups,
|
| dropout=dropout, with_attn=use_attn, size=now_res))
|
| pre_channel = channel_mult
|
| if not is_last:
|
| ups.append(Upsample(pre_channel))
|
| now_res = now_res * 2
|
| self.ups = nn.ModuleList(ups)
|
|
|
| self.final_conv = Block(pre_channel, default(out_channel, in_channel), groups=norm_groups)
|
| self.condition = CPEN(inchannel=condition_ch)
|
| self.condition_ch = condition_ch
|
|
|
| def forward(self, x, time, img_s1=None, class_label=None, return_condition=False, t_ori=0):
|
| condition = x[:, :self.condition_ch, ...].clone()
|
| x = x[:, self.condition_ch:, ...]
|
|
|
| c1, c2, c3, c4, c5 = self.condition(condition)
|
| c_base = [c1, c2, c3, c4, c5]
|
|
|
| c = []
|
| for i in range(len(c_base)):
|
| for _ in range(self.res_blocks):
|
| c.append(c_base[i])
|
|
|
| t = self.noise_level_mlp(time) if exists(self.noise_level_mlp) else None
|
|
|
| feats = []
|
| i = 0
|
| for layer in self.downs:
|
| if isinstance(layer, ResnetBlocWithAttn):
|
| x = layer(x, t, c[i])
|
| i += 1
|
| else:
|
| x = layer(x)
|
| feats.append(x)
|
|
|
| for layer in self.mid:
|
| if isinstance(layer, ResnetBlocWithAttn):
|
| x = layer(x, t, c5)
|
| else:
|
| x = layer(x)
|
|
|
| c_base = [c5, c4, c3, c2, c1]
|
| c = []
|
| for i in range(len(c_base)):
|
| for _ in range(self.res_blocks + 1):
|
| c.append(c_base[i])
|
|
|
| i = 0
|
| for layer in self.ups:
|
| if isinstance(layer, ResnetBlocWithAttn):
|
| x = layer(torch.cat((x, feats.pop()), dim=1), t, c[i])
|
| i += 1
|
| else:
|
| x = layer(x)
|
|
|
| if not return_condition:
|
| return self.final_conv(x)
|
| else:
|
| return self.final_conv(x), [c1, c2, c3, c4, c5]
|
|
|
|
|
|
|
|
|
|
|
|
|
| def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2):
|
| if schedule == 'linear':
|
| betas = np.linspace(linear_start, linear_end, n_timestep, dtype=np.float64)
|
| else:
|
| raise NotImplementedError(schedule)
|
| return betas
|
|
|
|
|
| class GaussianDiffusion(nn.Module):
|
| def __init__(self, denoise_fn, image_size, channels=3, schedule_opt=None, opt=None):
|
| super().__init__()
|
| self.channels = channels
|
| self.image_size = image_size
|
| self.denoise_fn = denoise_fn
|
| self.opt = opt
|
| self.ddim = schedule_opt.get('ddim', 1) if schedule_opt else 1
|
|
|
| def set_new_noise_schedule(self, schedule_opt, device, num_train_timesteps=1000):
|
| self.ddim = schedule_opt['ddim']
|
| self.num_train_timesteps = num_train_timesteps
|
| to_torch = partial(torch.tensor, dtype=torch.float32, device=device)
|
|
|
| betas = make_beta_schedule(
|
| schedule=schedule_opt['schedule'],
|
| n_timestep=num_train_timesteps,
|
| linear_start=schedule_opt['linear_start'],
|
| linear_end=schedule_opt['linear_end']
|
| )
|
|
|
| alphas = 1. - betas
|
| alphas_cumprod = np.cumprod(alphas, axis=0)
|
| self.sqrt_alphas_cumprod_prev = np.sqrt(np.append(1., alphas_cumprod))
|
|
|
| self.num_timesteps = int(betas.shape[0])
|
| self.register_buffer('betas', to_torch(betas))
|
| self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
|
|
|
| self.ddim_num_steps = schedule_opt['n_timestep']
|
| print(f'DDIM sampling steps: {self.ddim_num_steps}')
|
|
|
| def ddim_sample(self, condition_x, img_or_shape, device, seed=1):
|
| """DDIM sampling - matches the original E3Diff implementation."""
|
| eta = 0.8
|
|
|
| batch = img_or_shape[0]
|
| total_timesteps = self.num_train_timesteps
|
| sampling_timesteps = self.ddim_num_steps
|
|
|
| ts = torch.linspace(total_timesteps, 0, sampling_timesteps + 1).to(device).long()
|
| x = torch.randn(img_or_shape, device=device)
|
| batch_size = x.shape[0]
|
|
|
| imgs = [x]
|
| img_onestep = [condition_x[:, :self.channels, ...]]
|
|
|
| for i in range(1, sampling_timesteps + 1):
|
| cur_t = ts[i - 1] - 1
|
| prev_t = ts[i] - 1
|
|
|
| noise_level = torch.FloatTensor(
|
| [self.sqrt_alphas_cumprod_prev[cur_t.item()]]
|
| ).repeat(batch_size, 1).to(device)
|
|
|
| alpha_prod_t = self.alphas_cumprod[cur_t]
|
| alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else torch.tensor(1.0, device=device)
|
| beta_prod_t = 1 - alpha_prod_t
|
|
|
|
|
| model_output = self.denoise_fn(torch.cat([condition_x, x], dim=1), noise_level)
|
|
|
|
|
| sigma_2 = eta * (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
|
| noise = torch.randn_like(x)
|
|
|
|
|
| pred_original_sample = (x - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
|
| pred_original_sample = pred_original_sample.clamp(-1, 1)
|
|
|
| pred_sample_direction = (1 - alpha_prod_t_prev - sigma_2) ** 0.5 * model_output
|
|
|
| x = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction + sigma_2 ** 0.5 * noise
|
|
|
| imgs.append(x)
|
| img_onestep.append(pred_original_sample)
|
|
|
| imgs = torch.cat(imgs, dim=0)
|
| img_onestep = torch.cat(img_onestep, dim=0)
|
|
|
| return imgs, img_onestep
|
|
|
| @torch.no_grad()
|
| def super_resolution(self, x_in, continous=False, seed=1, img_s1=None):
|
| """Main inference method."""
|
| device = self.betas.device
|
| x = x_in
|
| shape = (x.shape[0], self.channels, x.shape[-2], x.shape[-1])
|
|
|
| self.ddim_num_steps = self.opt['ddim_steps']
|
| ret_img, img_onestep = self.ddim_sample(condition_x=x, img_or_shape=shape, device=device, seed=seed)
|
|
|
| if continous:
|
| return ret_img, img_onestep
|
| else:
|
| return ret_img[-x_in.shape[0]:], img_onestep
|
|
|
|
|
|
|
|
|
|
|
|
|
| class E3DiffInference:
|
| def __init__(self, weights_path=None, device="cuda", num_inference_steps=1):
|
| self.device = torch.device(device if torch.cuda.is_available() else "cpu")
|
| self.image_size = 256
|
| self.num_inference_steps = num_inference_steps
|
|
|
| print(f"[E3Diff] Initializing on device: {self.device}")
|
| print(f"[E3Diff] Inference steps: {num_inference_steps}")
|
|
|
| self.model = self._build_model()
|
| self._load_weights(weights_path)
|
| self.model.eval()
|
| print("[E3Diff] Model ready!")
|
|
|
| def _build_model(self):
|
| unet = UNet(
|
| in_channel=3,
|
| out_channel=3,
|
| norm_groups=16,
|
| inner_channel=64,
|
| channel_mults=[1, 2, 4, 8, 16],
|
| attn_res=[],
|
| res_blocks=1,
|
| dropout=0,
|
| image_size=self.image_size,
|
| condition_ch=3
|
| )
|
|
|
| schedule_opt = {
|
| 'schedule': 'linear',
|
| 'n_timestep': self.num_inference_steps,
|
| 'linear_start': 1e-6,
|
| 'linear_end': 1e-2,
|
| 'ddim': 1,
|
| 'lq_noiselevel': 0
|
| }
|
|
|
| opt = {
|
| 'stage': 2,
|
| 'ddim_steps': self.num_inference_steps,
|
| }
|
|
|
| model = GaussianDiffusion(
|
| denoise_fn=unet,
|
| image_size=self.image_size,
|
| channels=3,
|
| schedule_opt=schedule_opt,
|
| opt=opt
|
| )
|
|
|
| return model.to(self.device)
|
|
|
| def _load_weights(self, weights_path):
|
| if weights_path is None:
|
| weights_path = hf_hub_download(
|
| repo_id="Dhenenjay/E3Diff-SAR2Optical",
|
| filename="I700000_E719_gen.pth"
|
| )
|
|
|
| print(f"[E3Diff] Loading weights from: {weights_path}")
|
| state_dict = torch.load(weights_path, map_location=self.device, weights_only=False)
|
| self.model.load_state_dict(state_dict, strict=False)
|
| print("[E3Diff] Weights loaded!")
|
|
|
| def preprocess(self, image):
|
| if image.mode != 'RGB':
|
| image = image.convert('RGB')
|
| if image.size != (self.image_size, self.image_size):
|
| image = image.resize((self.image_size, self.image_size), Image.LANCZOS)
|
|
|
| img_np = np.array(image).astype(np.float32) / 255.0
|
| img_tensor = torch.from_numpy(img_np).permute(2, 0, 1)
|
| img_tensor = img_tensor * 2.0 - 1.0
|
| return img_tensor.unsqueeze(0).to(self.device)
|
|
|
| def postprocess(self, tensor):
|
| tensor = tensor.squeeze(0).cpu()
|
| tensor = torch.clamp(tensor, -1, 1)
|
| tensor = (tensor + 1.0) / 2.0
|
| img_np = (tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
|
| return Image.fromarray(img_np)
|
|
|
| @torch.no_grad()
|
| def translate(self, sar_image, seed=42):
|
| if seed is not None:
|
| torch.manual_seed(seed)
|
| np.random.seed(seed)
|
|
|
| sar_tensor = self.preprocess(sar_image)
|
|
|
| self.model.set_new_noise_schedule(
|
| {
|
| 'schedule': 'linear',
|
| 'n_timestep': self.num_inference_steps,
|
| 'linear_start': 1e-6,
|
| 'linear_end': 1e-2,
|
| 'ddim': 1,
|
| 'lq_noiselevel': 0
|
| },
|
| self.device,
|
| num_train_timesteps=1000
|
| )
|
|
|
| output, _ = self.model.super_resolution(sar_tensor, continous=False, seed=seed, img_s1=sar_tensor)
|
| return self.postprocess(output)
|
|
|
|
|
|
|
|
|
|
|
|
|
| class HighResProcessor:
|
| def __init__(self, device="cuda"):
|
| self.device = device
|
| self.model = None
|
| self.tile_size = 256
|
|
|
| def load_model(self, num_steps=1):
|
| print("Loading E3Diff model...")
|
| self.model = E3DiffInference(device=self.device, num_inference_steps=num_steps)
|
| self.num_steps = num_steps
|
|
|
| def create_blend_weights(self, tile_size, overlap):
|
| ramp = np.linspace(0, 1, overlap)
|
| weight = np.ones((tile_size, tile_size))
|
| weight[:overlap, :] *= ramp[:, np.newaxis]
|
| weight[-overlap:, :] *= ramp[::-1, np.newaxis]
|
| weight[:, :overlap] *= ramp[np.newaxis, :]
|
| weight[:, -overlap:] *= ramp[np.newaxis, ::-1]
|
| return weight[:, :, np.newaxis]
|
|
|
| def process(self, image, overlap=64, num_steps=1):
|
| if self.model is None or self.num_steps != num_steps:
|
| self.load_model(num_steps)
|
|
|
| if isinstance(image, Image.Image):
|
| if image.mode != 'RGB':
|
| image = image.convert('RGB')
|
| img_np = np.array(image).astype(np.float32) / 255.0
|
| else:
|
| img_np = image
|
|
|
| h, w = img_np.shape[:2]
|
| tile_size = self.tile_size
|
| step = tile_size - overlap
|
|
|
| pad_h = (step - (h - overlap) % step) % step
|
| pad_w = (step - (w - overlap) % step) % step
|
| img_padded = np.pad(img_np, ((0, pad_h), (0, pad_w), (0, 0)), mode='reflect')
|
|
|
| h_pad, w_pad = img_padded.shape[:2]
|
|
|
| output = np.zeros((h_pad, w_pad, 3), dtype=np.float32)
|
| weights = np.zeros((h_pad, w_pad, 1), dtype=np.float32)
|
| blend_weight = self.create_blend_weights(tile_size, overlap)
|
|
|
| y_positions = list(range(0, h_pad - tile_size + 1, step))
|
| x_positions = list(range(0, w_pad - tile_size + 1, step))
|
| total_tiles = len(y_positions) * len(x_positions)
|
|
|
| print(f"Processing {total_tiles} tiles at {w}x{h}...")
|
|
|
| tile_idx = 0
|
| for y in y_positions:
|
| for x in x_positions:
|
| tile = img_padded[y:y+tile_size, x:x+tile_size]
|
| tile_pil = Image.fromarray((tile * 255).astype(np.uint8))
|
|
|
| result_pil = self.model.translate(tile_pil, seed=42)
|
| result = np.array(result_pil).astype(np.float32) / 255.0
|
|
|
| output[y:y+tile_size, x:x+tile_size] += result * blend_weight
|
| weights[y:y+tile_size, x:x+tile_size] += blend_weight
|
|
|
| tile_idx += 1
|
| print(f" Tile {tile_idx}/{total_tiles}")
|
|
|
| output = output / (weights + 1e-8)
|
| output = output[:h, :w]
|
|
|
| return (output * 255).astype(np.uint8)
|
|
|
| def enhance(self, image, contrast=1.1, sharpness=1.15, color=1.1):
|
| if isinstance(image, np.ndarray):
|
| image = Image.fromarray(image)
|
| image = ImageEnhance.Contrast(image).enhance(contrast)
|
| image = ImageEnhance.Sharpness(image).enhance(sharpness)
|
| image = ImageEnhance.Color(image).enhance(color)
|
| return image
|
|
|
|
|
|
|
|
|
|
|
|
|
| processor = None
|
|
|
| def load_sar_image(filepath):
|
| """Load SAR image from various formats."""
|
| try:
|
| import rasterio
|
| with rasterio.open(filepath) as src:
|
| data = src.read(1)
|
| if data.dtype in [np.float32, np.float64]:
|
| valid = data[np.isfinite(data)]
|
| if len(valid) > 0:
|
| p2, p98 = np.percentile(valid, [2, 98])
|
| data = np.clip(data, p2, p98)
|
| data = ((data - p2) / (p98 - p2 + 1e-8) * 255).astype(np.uint8)
|
| elif data.dtype == np.uint16:
|
| p2, p98 = np.percentile(data, [2, 98])
|
| data = np.clip(data, p2, p98)
|
| data = ((data - p2) / (p98 - p2 + 1e-8) * 255).astype(np.uint8)
|
| return Image.fromarray(data).convert('RGB')
|
| except:
|
| pass
|
|
|
| return Image.open(filepath).convert('RGB')
|
|
|
|
|
| def _translate_sar_impl(file, num_steps, overlap, enhance_output):
|
| """Main translation function implementation."""
|
| global processor
|
|
|
| if file is None:
|
| return None, None, "Please upload a SAR image"
|
|
|
| if processor is None:
|
| processor = HighResProcessor()
|
|
|
| print("Processing SAR image...")
|
|
|
| filepath = file.name if hasattr(file, 'name') else file
|
| image = load_sar_image(filepath)
|
|
|
| w, h = image.size
|
| print(f"Input size: {w}x{h}")
|
|
|
| start = time.time()
|
| result = processor.process(image, overlap=int(overlap), num_steps=int(num_steps))
|
| elapsed = time.time() - start
|
|
|
| result_pil = Image.fromarray(result)
|
|
|
| if enhance_output:
|
| result_pil = processor.enhance(result_pil)
|
|
|
| tiff_path = tempfile.mktemp(suffix='.tiff')
|
| result_pil.save(tiff_path, format='TIFF', compression='lzw')
|
|
|
| print(f"Complete in {elapsed:.1f}s!")
|
|
|
| info = f"Processed in {elapsed:.1f}s | Output: {result_pil.size[0]}x{result_pil.size[1]}"
|
|
|
| return result_pil, tiff_path, info
|
|
|
|
|
|
|
| if GPU_AVAILABLE and spaces is not None:
|
| translate_sar = spaces.GPU(duration=300)(_translate_sar_impl)
|
| else:
|
| translate_sar = _translate_sar_impl
|
|
|
|
|
|
|
| with gr.Blocks(title="E3Diff: SAR-to-Optical Translation") as demo:
|
| gr.Markdown("""
|
| # 🛰️ E3Diff: High-Resolution SAR-to-Optical Translation
|
|
|
| **CVPR PBVS2025 Challenge Winner** | Upload any SAR image and get a photorealistic optical translation.
|
|
|
| - Supports full resolution processing with seamless tiling
|
| - Multiple quality levels (1-8 inference steps)
|
| - TIFF output for commercial use
|
| """)
|
|
|
| with gr.Row():
|
| with gr.Column():
|
| input_file = gr.File(label="SAR Input (TIFF, PNG, JPG supported)", file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"])
|
|
|
| with gr.Row():
|
| num_steps = gr.Slider(1, 8, value=1, step=1, label="Quality Steps (1=fast, 8=best)")
|
| overlap = gr.Slider(16, 128, value=64, step=16, label="Tile Overlap")
|
|
|
| enhance = gr.Checkbox(value=True, label="Apply enhancement")
|
|
|
| submit_btn = gr.Button("🚀 Translate to Optical", variant="primary")
|
|
|
| with gr.Column():
|
| output_image = gr.Image(label="Optical Output")
|
| output_file = gr.File(label="Download TIFF")
|
| info_text = gr.Textbox(label="Processing Info")
|
|
|
| submit_btn.click(
|
| fn=translate_sar,
|
| inputs=[input_file, num_steps, overlap, enhance],
|
| outputs=[output_image, output_file, info_text]
|
| )
|
|
|
| gr.Markdown("""
|
| ---
|
| **Tips:** The model works best with Sentinel-1 style SAR imagery. Use steps=1 for speed, steps=4-8 for quality.
|
| """)
|
|
|
|
|
| if __name__ == "__main__":
|
| demo.queue().launch(ssr_mode=False)
|
|
|
