Add StarNet PyTorch source code

StarNetPyTorch/convert_tf2_h5.py

@@ -0,0 +1,136 @@
+import argparse
+from pathlib import Path
+
+import h5py
+import numpy as np
+import torch
+
+from model import Generator
+
+
+def _decode_names(values):
+    return [value.decode("utf-8") if isinstance(value, bytes) else str(value) for value in values]
+
+
+def _read_layer_weights(layer_group):
+    weight_names = _decode_names(layer_group.attrs.get("weight_names", []))
+    weights = {}
+    for weight_name in weight_names:
+        short_name = weight_name.rsplit("/", 1)[-1].split(":", 1)[0]
+        try:
+            weights[short_name] = np.asarray(layer_group[weight_name])
+        except KeyError:
+            relative_name = weight_name.split("/", 1)[-1]
+            weights[short_name] = np.asarray(layer_group[relative_name])
+    return weights
+
+
+def _load_keras_layers(h5_path):
+    layers = []
+    with h5py.File(h5_path, "r") as handle:
+        root = handle
+        if "layer_names" not in root.attrs and "model_weights" in root:
+            root = root["model_weights"]
+
+        layer_names = _decode_names(root.attrs.get("layer_names", []))
+        if not layer_names:
+            raise ValueError(
+                "Unsupported H5 layout: no Keras layer_names attribute was found. "
+                "Use a TensorFlow/Keras save_weights H5 file."
+            )
+
+        for layer_name in layer_names:
+            layer_group = root[layer_name]
+            weights = _read_layer_weights(layer_group)
+            if weights:
+                layers.append((layer_name, weights))
+    return layers
+
+
+def _copy_conv(state, torch_name, weights, transpose=False):
+    required = {"kernel", "bias"}
+    if not required.issubset(weights):
+        raise ValueError("Layer %s is missing kernel or bias" % torch_name)
+
+    kernel = torch.from_numpy(weights["kernel"]).permute(3, 2, 0, 1).contiguous()
+    if transpose:
+        state[torch_name + ".weight"] = kernel
+    else:
+        state[torch_name + ".weight"] = kernel
+    state[torch_name + ".bias"] = torch.from_numpy(weights["bias"])
+
+
+def _copy_batch_norm(state, torch_name, weights):
+    names = {
+        "weight": "gamma",
+        "bias": "beta",
+        "running_mean": "moving_mean",
+        "running_var": "moving_variance",
+    }
+    missing = [source for source in names.values() if source not in weights]
+    if missing:
+        raise ValueError("Layer %s is missing: %s" % (torch_name, ", ".join(missing)))
+    for target, source in names.items():
+        state[torch_name + "." + target] = torch.from_numpy(weights[source])
+
+
+def convert(input_path, output_path, mode):
+    input_channels = 3 if mode == "RGB" else 1
+    layers = _load_keras_layers(input_path)
+
+    conv_layers = [weights for name, weights in layers if name.startswith("conv2d") and "transpose" not in name]
+    deconv_layers = [weights for name, weights in layers if name.startswith("conv2d_transpose")]
+    batch_norm_layers = [weights for name, weights in layers if name.startswith("batch_normalization")]
+
+    if len(conv_layers) != 8 or len(deconv_layers) != 8 or len(batch_norm_layers) != 14:
+        raise ValueError(
+            "Unexpected generator layout: found %d Conv2D, %d Conv2DTranspose, and "
+            "%d BatchNormalization layers; expected 8, 8, and 14."
+            % (len(conv_layers), len(deconv_layers), len(batch_norm_layers))
+        )
+
+    generator = Generator(input_channels=input_channels, batch_norm_eps=1e-3)
+    state = generator.state_dict()
+
+    for index, weights in enumerate(conv_layers):
+        _copy_conv(state, "down_convs.%d" % index, weights)
+    for index, weights in enumerate(deconv_layers):
+        _copy_conv(state, "up_convs.%d" % index, weights, transpose=True)
+    for index, weights in enumerate(batch_norm_layers[:7]):
+        _copy_batch_norm(state, "down_bns.%d" % index, weights)
+    for index, weights in enumerate(batch_norm_layers[7:]):
+        _copy_batch_norm(state, "up_bns.%d" % index, weights)
+
+    generator.load_state_dict(state)
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "state_dict": generator.state_dict(),
+            "mode": mode,
+            "source": str(Path(input_path)),
+        },
+        output_path,
+    )
+    return output_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert StarNet TF2 Keras generator H5 weights")
+    parser.add_argument("input", help="Keras generator H5 file, such as weights_G_Greyscale.h5")
+    parser.add_argument(
+        "output",
+        nargs="?",
+        default=None,
+        help="Output .pth file; defaults to the input filename with a .pth suffix",
+    )
+    parser.add_argument("--mode", choices=("RGB", "Greyscale"), default="Greyscale")
+    args = parser.parse_args()
+
+    output = args.output or str(Path(args.input).with_suffix(".pth"))
+    output_path = convert(args.input, output, args.mode)
+    print("Saved PyTorch weights: %s" % output_path)
+
+
+if __name__ == "__main__":
+    main()

StarNetPyTorch/convert_tf_checkpoint.py

@@ -0,0 +1,80 @@
+import argparse
+from pathlib import Path
+
+import torch
+
+from model import Generator
+
+
+def _read_tf_checkpoint(checkpoint_path):
+    try:
+        import tensorflow as tf
+    except ImportError as exc:
+        raise RuntimeError("TensorFlow is required only for checkpoint conversion.") from exc
+
+    return tf.train.load_checkpoint(str(checkpoint_path))
+
+
+def _copy_conv(reader, state, tf_name, torch_name):
+    state[torch_name + ".weight"] = torch.from_numpy(reader.get_tensor(tf_name + "/kernel")).permute(3, 2, 0, 1)
+    state[torch_name + ".bias"] = torch.from_numpy(reader.get_tensor(tf_name + "/bias"))
+
+
+def _copy_deconv(reader, state, tf_name, torch_name):
+    state[torch_name + ".weight"] = torch.from_numpy(reader.get_tensor(tf_name + "/kernel")).permute(3, 2, 0, 1)
+    state[torch_name + ".bias"] = torch.from_numpy(reader.get_tensor(tf_name + "/bias"))
+
+
+def _copy_bn(reader, state, tf_name, torch_name):
+    state[torch_name + ".weight"] = torch.from_numpy(reader.get_tensor(tf_name + "/gamma"))
+    state[torch_name + ".bias"] = torch.from_numpy(reader.get_tensor(tf_name + "/beta"))
+    state[torch_name + ".running_mean"] = torch.from_numpy(reader.get_tensor(tf_name + "/moving_mean"))
+    state[torch_name + ".running_var"] = torch.from_numpy(reader.get_tensor(tf_name + "/moving_variance"))
+
+
+def convert(checkpoint_path, output_path):
+    checkpoint_path = Path(checkpoint_path)
+    output_path = Path(output_path)
+
+    reader = _read_tf_checkpoint(checkpoint_path)
+    generator = Generator()
+    state = generator.state_dict()
+
+    _copy_conv(reader, state, "generator/g_conv0/conv2d", "down_convs.0")
+    for i in range(1, 8):
+        _copy_conv(reader, state, "generator/g_conv%d/conv2d" % i, "down_convs.%d" % i)
+        _copy_bn(
+            reader,
+            state,
+            "generator/g_conv%d/batch_normalization" % i,
+            "down_bns.%d" % (i - 1),
+        )
+
+    for i in range(8):
+        _copy_deconv(reader, state, "generator/g_deconv%d/conv2d_transpose" % i, "up_convs.%d" % i)
+        if i < 7:
+            _copy_bn(
+                reader,
+                state,
+                "generator/g_deconv%d/batch_normalization" % i,
+                "up_bns.%d" % i,
+            )
+
+    generator.load_state_dict(state)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"state_dict": generator.state_dict()}, output_path)
+    return output_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert StarNet TF1 generator checkpoint to PyTorch")
+    parser.add_argument("--checkpoint", default="model.ckpt", help="TensorFlow checkpoint prefix")
+    parser.add_argument("--output", default="starnet_generator.pth", help="Output PyTorch weights")
+    args = parser.parse_args()
+
+    output_path = convert(args.checkpoint, args.output)
+    print("Saved PyTorch weights to %s" % output_path)
+
+
+if __name__ == "__main__":
+    main()

StarNetPyTorch/cvt.sh

@@ -0,0 +1,13 @@
+python convert_tf2_h5.py \
+  /data/xch/chm_workspace2/Models/starnet_weights2/weights_D_Greyscale.h5 \
+  /data/xch/chm_workspace2/Models/starnet_weights2/weights_D_Greyscale.pth \
+  --mode Greyscale
+
+
+
+  # /data/xch/chm_workspace2/Models/starnet_weights2/weights_G_RGB.h5 \
+  # /data/xch/chm_workspace2/Models/starnet_weights2/weights_G_RGB.pth \
+
+#   /data/xch/chm_workspace2/Models/starnet_weights2/weights_G_Greyscale.h5 \
+#   /data/xch/chm_workspace2/Models/starnet_weights2/weights_G_Greyscale.pth \
+#   --mode Greyscale

StarNetPyTorch/model.py

@@ -0,0 +1,190 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TFBatchNorm2d(nn.Module):
+    """BatchNorm compatible with the original TF graph's training=True usage."""
+
+    def __init__(self, channels, eps=1e-5, momentum=0.99):
+        super().__init__()
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = nn.Parameter(torch.ones(channels))
+        self.bias = nn.Parameter(torch.zeros(channels))
+        self.register_buffer("running_mean", torch.zeros(channels))
+        self.register_buffer("running_var", torch.ones(channels))
+
+    def forward(self, x):
+        mean = x.mean(dim=(0, 2, 3), keepdim=True)
+        var = (x - mean).pow(2).mean(dim=(0, 2, 3), keepdim=True)
+        if self.training:
+            with torch.no_grad():
+                self.running_mean.mul_(self.momentum).add_(
+                    mean.reshape(-1), alpha=1.0 - self.momentum
+                )
+                self.running_var.mul_(self.momentum).add_(
+                    var.reshape(-1), alpha=1.0 - self.momentum
+                )
+        weight = self.weight.view(1, -1, 1, 1)
+        bias = self.bias.view(1, -1, 1, 1)
+        return (x - mean) * torch.rsqrt(var + self.eps) * weight + bias
+
+
+class Generator(nn.Module):
+    def __init__(self, input_channels=3, batch_norm_eps=1e-5):
+        super().__init__()
+        filters = [64, 128, 256, 512, 512, 512, 512, 512, 512, 512, 512, 512, 256, 128, 64]
+
+        self.down_convs = nn.ModuleList()
+        in_channels = input_channels
+        for out_channels in filters[:8]:
+            self.down_convs.append(
+                nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
+            )
+            in_channels = out_channels
+
+        self.down_bns = nn.ModuleList(
+            [TFBatchNorm2d(channels, eps=batch_norm_eps) for channels in filters[1:8]]
+        )
+
+        self.up_convs = nn.ModuleList()
+        up_in_channels = [
+            filters[7],
+            filters[8] + filters[6],
+            filters[9] + filters[5],
+            filters[10] + filters[4],
+            filters[11] + filters[3],
+            filters[12] + filters[2],
+            filters[13] + filters[1],
+            filters[14] + filters[0],
+        ]
+        up_out_channels = filters[8:] + [input_channels]
+        for in_channels, out_channels in zip(up_in_channels, up_out_channels):
+            self.up_convs.append(
+                nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
+            )
+
+        self.up_bns = nn.ModuleList(
+            [TFBatchNorm2d(channels, eps=batch_norm_eps) for channels in filters[8:15]]
+        )
+        self.apply(_initialize_layer)
+
+    def forward(self, x):
+        original = x
+        layers = []
+
+        y = self.down_convs[0](x)
+        layers.append(y)
+
+        for i in range(1, 8):
+            y = F.leaky_relu(layers[-1], negative_slope=0.2)
+            y = self.down_convs[i](y)
+            y = self.down_bns[i - 1](y)
+            layers.append(y)
+
+        y = F.relu(layers[-1])
+        y = self.up_convs[0](y)
+        y = self.up_bns[0](y)
+        layers.append(y)
+
+        for up_index, skip_index in enumerate([6, 5, 4, 3, 2, 1], start=1):
+            y = torch.cat([layers[-1], layers[skip_index]], dim=1)
+            y = F.relu(y)
+            y = self.up_convs[up_index](y)
+            y = self.up_bns[up_index](y)
+            layers.append(y)
+
+        y = torch.cat([layers[-1], layers[0]], dim=1)
+        y = F.relu(y)
+        y = self.up_convs[7](y)
+        y = F.relu(y)
+        return original - y
+
+
+class Discriminator(nn.Module):
+    def __init__(self, input_channels=3, batch_norm_eps=1e-3):
+        super().__init__()
+        filters = [32, 64, 64, 128, 128, 256, 256, 256, 8]
+        strides = [1, 2, 1, 2, 1, 2, 1, 2, 2]
+
+        self.convs = nn.ModuleList()
+        in_channels = input_channels
+        for index, (out_channels, stride) in enumerate(zip(filters, strides)):
+            padding = 1 if stride == 1 else 0
+            self.convs.append(
+                nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=3,
+                    stride=stride,
+                    padding=padding,
+                )
+            )
+            in_channels = out_channels
+
+        self.batch_norms = nn.ModuleList(
+            [TFBatchNorm2d(channels, eps=batch_norm_eps) for channels in filters[1:]]
+        )
+        self.dense = nn.Conv2d(filters[-1], 1, kernel_size=1)
+        self.apply(_initialize_layer)
+
+    def forward(self, x):
+        features = []
+
+        y = F.leaky_relu(self.convs[0](x), negative_slope=0.2)
+        features.append(y)
+
+        for index in range(1, len(self.convs)):
+            y = self.convs[index](features[-1])
+            y = self.batch_norms[index - 1](y)
+            y = F.leaky_relu(y, negative_slope=0.2)
+            features.append(y)
+
+        prediction = torch.sigmoid(self.dense(features[-1]))
+        return (*features[:8], prediction)
+
+
+def _initialize_layer(module):
+    if isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
+        nn.init.xavier_uniform_(module.weight)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def create_generator(
+    weights_path=None,
+    device=None,
+    input_channels=3,
+    batch_norm_eps=1e-5,
+):
+    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net = Generator(
+        input_channels=input_channels,
+        batch_norm_eps=batch_norm_eps,
+    ).to(device)
+    if weights_path:
+        checkpoint = torch.load(weights_path, map_location=device)
+        state_dict = checkpoint.get("state_dict", checkpoint)
+        net.load_state_dict(state_dict)
+    net.eval()
+    return net
+
+
+def create_discriminator(
+    weights_path=None,
+    device=None,
+    input_channels=3,
+    batch_norm_eps=1e-3,
+):
+    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net = Discriminator(
+        input_channels=input_channels,
+        batch_norm_eps=batch_norm_eps,
+    ).to(device)
+    if weights_path:
+        checkpoint = torch.load(weights_path, map_location=device)
+        state_dict = checkpoint.get("state_dict", checkpoint)
+        net.load_state_dict(state_dict)
+    net.eval()
+    return net

StarNetPyTorch/requirements.txt

@@ -0,0 +1,5 @@
+numpy>=1.24
+tifffile>=2023.7
+torch>=2.1
+matplotlib>=3.7
+

StarNetPyTorch/run

@@ -0,0 +1,2 @@
+
+

StarNetPyTorch/run.sh

@@ -0,0 +1,11 @@
+python starnet_v1_pytorch_transform.py \
+  rgb_test5.tif \
+  rgb_test5_starless.tif \
+  --weights ../starnet_weights2/weights \
+  --mode RGB \
+  --window-size 512 \
+  --stride 128 \
+  --device cuda
+
+
+  

StarNetPyTorch/starnet.py

@@ -0,0 +1,49 @@
+import argparse
+import time
+
+import transform as transform_module
+
+
+epochs = 10000
+batch = 1
+steps = 1000
+output_freq = 100
+verbose = True
+images = True
+log_freq = 50
+gen_plots = True
+learning_rates = [0.000002, 0.0000005]
+stride = 64
+
+
+def main():
+    parser = argparse.ArgumentParser(description="PyTorch StarNet transform entry point")
+    subparsers = parser.add_subparsers(dest="command")
+
+    transform_parser = subparsers.add_parser("transform", help="Remove stars from a TIFF image")
+    transform_parser.add_argument("image", help="Input TIFF image")
+    transform_parser.add_argument("--stride", type=int, default=stride, help="Sliding window stride")
+    transform_parser.add_argument("--weights", default=None, help="Path to starnet_generator.pth")
+    transform_parser.add_argument("--device", default=None, help="Torch device, for example cuda or cpu")
+
+    args = parser.parse_args()
+    if args.command == "transform":
+        start = time.time()
+        transform_module.transform(
+            imageName=args.image,
+            stride=args.stride,
+            weights_path=args.weights,
+            device=args.device,
+        )
+        stop = time.time()
+        minutes = float((stop - start) / 60)
+        if minutes > 60.0:
+            print("Total time taken: %.1f hours" % (minutes / 60))
+        else:
+            print("Total time taken: %.1f minutes" % minutes)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

StarNetPyTorch/starnet_v1_pytorch.py

@@ -0,0 +1,554 @@
+import pickle
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import tifffile as tiff
+import torch
+import torch.nn.functional as F
+
+from model import Discriminator, Generator
+
+
+class StarNet:
+    def __init__(
+        self,
+        mode,
+        window_size=512,
+        stride=256,
+        lr=1e-4,
+        train_folder="./train/",
+        batch_size=1,
+        device=None,
+    ):
+        if mode not in ("RGB", "Greyscale"):
+            raise ValueError("Mode should be either RGB or Greyscale")
+        if window_size % 256 != 0:
+            raise ValueError("window_size must be divisible by 256")
+        if stride <= 0 or stride > window_size:
+            raise ValueError("stride must be greater than zero and no larger than window_size")
+        if (window_size - stride) % 2 != 0:
+            raise ValueError("window_size - stride must be even")
+
+        self.mode = mode
+        self.input_channels = 3 if mode == "RGB" else 1
+        self.window_size = window_size
+        self.stride = stride
+        self.lr = lr
+        self.train_folder = Path(train_folder)
+        self.batch_size = batch_size
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+
+        self.history = {}
+        self._ema = 0.9999
+        self.weights = []
+        self.original = []
+        self.starless = []
+        self.iters_per_epoch = 0
+
+        self.G = None
+        self.D = None
+        self.gen_optimizer = None
+        self.dis_optimizer = None
+
+    def __str__(self):
+        return "StarNet PyTorch instance"
+
+    def load_training_dataset(self):
+        original_dir = self.train_folder / "original"
+        starless_dir = self.train_folder / "starless"
+        original_files = sorted(path.name for path in original_dir.glob("*.tif"))
+        starless_files = sorted(path.name for path in starless_dir.glob("*.tif"))
+
+        if not original_files or not starless_files:
+            raise ValueError("No training data found in %s" % self.train_folder)
+        if original_files != starless_files:
+            raise ValueError(
+                "Corresponding names in the original and starless folders must be equal"
+            )
+
+        print("Total training images found: %d" % len(original_files))
+        self.original = [tiff.imread(original_dir / name) for name in original_files]
+        self.starless = [tiff.imread(starless_dir / name) for name in starless_files]
+        self.weights = []
+        total_pixels = 0
+
+        for name, original, starless in zip(original_files, self.original, self.starless):
+            if original.shape != starless.shape:
+                raise ValueError("Image sizes are not equal for %s" % name)
+            if original.shape[0] < self.window_size or original.shape[1] < self.window_size:
+                raise ValueError(
+                    "%s is smaller than the configured window size %d"
+                    % (name, self.window_size)
+                )
+            pixels = original.shape[0] * original.shape[1]
+            total_pixels += pixels
+            self.weights.append(pixels)
+
+        self.iters_per_epoch = max(1, total_pixels // (self.window_size * self.window_size))
+        self.weights = np.asarray(self.weights, dtype=np.float64)
+        self.weights /= self.weights.sum()
+
+        print("Total size of training images: %.2f MP" % (total_pixels / 1e6))
+        print("One epoch is set to %d iterations" % self.iters_per_epoch)
+        print("Training dataset has been successfully loaded!")
+
+    def _generator(self, m=64):
+        if m != 64:
+            raise ValueError("The translated architecture supports m=64")
+        return Generator(input_channels=self.input_channels, batch_norm_eps=1e-3)
+
+    def _discriminator(self):
+        return Discriminator(input_channels=self.input_channels, batch_norm_eps=1e-3)
+
+    def _weight_paths(self, weights):
+        path = Path(weights)
+        
+        if path.suffix.lower() in (".pth", ".pt"):
+            return path, None
+        return (
+            Path(str(path) + "_G_" + self.mode + ".pth"),
+            Path(str(path) + "_D_" + self.mode + ".pth"),
+        )
+
+    @staticmethod
+    def _load_state(module, path, device):
+        try:
+            checkpoint = torch.load(path, map_location=device, weights_only=True)
+        except TypeError:  # PyTorch < 2.0 compatibility
+            checkpoint = torch.load(path, map_location=device)
+        state_dict = checkpoint.get("state_dict", checkpoint)
+        module.load_state_dict(state_dict)
+
+    def load_model(self, weights=None, history=None):
+        self.G = self._generator().to(self.device)
+        self.D = self._discriminator().to(self.device)
+        self.gen_optimizer = torch.optim.Adam(self.G.parameters(), lr=self.lr)
+        self.dis_optimizer = torch.optim.Adam(self.D.parameters(), lr=self.lr / 4)
+
+        if weights:
+            generator_path, discriminator_path = self._weight_paths(weights)
+            if not generator_path.exists():
+                raise FileNotFoundError("Generator weights were not found at %s" % generator_path)
+            self._load_state(self.G, generator_path, self.device)
+
+            if discriminator_path and discriminator_path.exists():
+                self._load_state(self.D, discriminator_path, self.device)
+            elif discriminator_path:
+                print("Discriminator weights were not found; using initialized weights.")
+
+        if history:
+            history_path = Path(str(history) + "_" + self.mode + ".pkl")
+            with history_path.open("rb") as handle:
+                self.history = pickle.load(handle)
+
+        self.G.eval()
+        self.D.eval()
+        return self.G
+
+    def initialize_model(self):
+        return self.load_model()
+
+    @staticmethod
+    def _ramp(x):
+        if torch.is_tensor(x):
+            return torch.clamp(x, 0.0, 1.0)
+        return np.clip(x, 0.0, 1.0)
+
+    @staticmethod
+    def _image_scale(image):
+        if image.dtype == np.uint8:
+            return 255.0
+        if image.dtype == np.uint16:
+            return 65535.0
+        if np.issubdtype(image.dtype, np.floating):
+            return 1.0
+        raise ValueError("Unsupported training image dtype: %s" % image.dtype)
+
+    def _get_sample(self, r, h, w, type):
+        if type not in ("original", "starless"):
+            raise ValueError("type must be original or starless")
+        source = self.original if type == "original" else self.starless
+        sample = source[r][h:h + self.window_size, w:w + self.window_size]
+        return sample.astype(np.float32) / self._image_scale(sample)
+
+    def _augmentator(self, original, starless):
+        original = np.array(original, copy=True)
+        starless = np.array(starless, copy=True)
+
+        if np.random.rand() < 0.5:
+            original = np.flip(original, axis=1)
+            starless = np.flip(starless, axis=1)
+        if np.random.rand() < 0.5:
+            original = np.flip(original, axis=0)
+            starless = np.flip(starless, axis=0)
+        if np.random.rand() < 0.5:
+            rotations = np.random.randint(1, 4)
+            original = np.rot90(original, rotations, axes=(1, 0))
+            starless = np.rot90(starless, rotations, axes=(1, 0))
+
+        if self.mode == "RGB":
+            if original.ndim != 3 or original.shape[2] < 3:
+                raise ValueError("RGB training images must contain at least three channels")
+            original = original[:, :, :3]
+            starless = starless[:, :, :3]
+            if np.random.rand() < 0.7:
+                channel = np.random.randint(3)
+                minimum = min(float(original.min()), float(starless.min()))
+                offset = np.random.rand() * 0.25 - np.random.rand() * minimum
+                original[:, :, channel] += offset * (1.0 - original[:, :, channel])
+                starless[:, :, channel] += offset * (1.0 - starless[:, :, channel])
+            if np.random.rand() < 0.7:
+                sequence = np.random.permutation(3)
+                original = original[:, :, sequence]
+                starless = starless[:, :, sequence]
+        else:
+            if np.random.rand() < 0.7:
+                minimum = min(float(original.min()), float(starless.min()))
+                offset = np.random.rand() * 0.25 - np.random.rand() * minimum
+                original += offset * (1.0 - original)
+                starless += offset * (1.0 - starless)
+            if original.ndim == 3:
+                channel = np.random.randint(min(3, original.shape[2]))
+                original = original[:, :, channel]
+                starless = starless[:, :, channel]
+            original = original[:, :, None]
+            starless = starless[:, :, None]
+
+        return (
+            np.ascontiguousarray(np.clip(original, 0.0, 1.0)),
+            np.ascontiguousarray(np.clip(starless, 0.0, 1.0)),
+        )
+
+    def generate_input(self, iterations=1, augmentation=False):
+        if not self.original:
+            raise RuntimeError("Training dataset was not loaded")
+
+        original_batch = None
+        starless_batch = None
+        for _ in range(iterations):
+            original_batch = np.zeros(
+                (
+                    self.batch_size,
+                    self.window_size,
+                    self.window_size,
+                    self.input_channels,
+                ),
+                dtype=np.float32,
+            )
+            starless_batch = np.zeros_like(original_batch)
+
+            for index in range(self.batch_size):
+                image_index = int(
+                    np.random.choice(len(self.original), p=np.asarray(self.weights))
+                )
+                image_h, image_w = self.original[image_index].shape[:2]
+                top = np.random.randint(0, image_h - self.window_size + 1)
+                left = np.random.randint(0, image_w - self.window_size + 1)
+                original = self._get_sample(image_index, top, left, "original")
+                starless = self._get_sample(image_index, top, left, "starless")
+
+                if augmentation:
+                    original, starless = self._augmentator(original, starless)
+                elif self.mode == "Greyscale":
+                    if original.ndim == 3:
+                        channel = np.random.randint(min(3, original.shape[2]))
+                        original = original[:, :, channel]
+                        starless = starless[:, :, channel]
+                    original = original[:, :, None]
+                    starless = starless[:, :, None]
+                else:
+                    original = original[:, :, :3]
+                    starless = starless[:, :, :3]
+
+                original_batch[index] = original
+                starless_batch[index] = starless
+
+        return original_batch, starless_batch
+
+    @staticmethod
+    def _set_requires_grad(module, enabled):
+        for parameter in module.parameters():
+            parameter.requires_grad_(enabled)
+
+    def _losses(self, original, target):
+        generated = self.G(original)
+        real_outputs = self.D(target)
+        fake_outputs = self.D(generated)
+
+        real_features = real_outputs[:-1]
+        fake_features = fake_outputs[:-1]
+        predict_real = real_outputs[-1]
+        predict_fake = fake_outputs[-1]
+
+        dis_loss = -(
+            torch.log(predict_real + 1e-8) + torch.log(1.0 - predict_fake + 1e-8)
+        ).mean()
+        gen_loss_gan = -torch.log(predict_fake + 1e-8).mean()
+        perceptual = [
+            torch.mean(torch.abs(fake - real))
+            for fake, real in zip(fake_features, real_features)
+        ]
+        gen_l1 = F.l1_loss(generated, target)
+        gen_loss = (
+            gen_loss_gan * 0.1
+            + perceptual[0] * 0.1
+            + sum(loss * 10.0 for loss in perceptual[1:])
+            + gen_l1 * 100.0
+        )
+
+        metrics = {
+            "dis_loss": dis_loss,
+            "gen_loss_GAN": gen_loss_gan,
+            "gen_p1": perceptual[0],
+            "gen_p2": perceptual[1],
+            "gen_p3": perceptual[2],
+            "gen_p4": perceptual[3],
+            "gen_p5": perceptual[4],
+            "gen_p6": perceptual[5],
+            "gen_p7": perceptual[6],
+            "gen_p8": perceptual[7],
+            "gen_L1": gen_l1 * 100.0,
+            "total": gen_loss,
+        }
+        return gen_loss, dis_loss, metrics
+
+    def _train_step(self, original, target):
+        self.G.train()
+        self.D.train()
+        self.gen_optimizer.zero_grad(set_to_none=True)
+        self.dis_optimizer.zero_grad(set_to_none=True)
+
+        gen_loss, dis_loss, metrics = self._losses(original, target)
+        generator_gradients = torch.autograd.grad(
+            gen_loss,
+            tuple(self.G.parameters()),
+            retain_graph=True,
+        )
+        discriminator_gradients = torch.autograd.grad(
+            dis_loss,
+            tuple(self.D.parameters()),
+        )
+
+        for parameter, gradient in zip(self.G.parameters(), generator_gradients):
+            parameter.grad = gradient
+        for parameter, gradient in zip(self.D.parameters(), discriminator_gradients):
+            parameter.grad = gradient
+
+        self.gen_optimizer.step()
+        self.dis_optimizer.step()
+        return {name: float(value.detach().cpu()) for name, value in metrics.items()}
+
+    def _update_history(self, metrics):
+        for name, value in metrics.items():
+            if name in self.history and self.history[name]:
+                value = value * (1.0 - self._ema) + self.history[name][-1] * self._ema
+                self.history[name].append(value)
+            else:
+                self.history[name] = [value]
+
+    def _plot_progress(self, original, target):
+        self.G.eval()
+        with torch.inference_mode():
+            generated = self.G(original)
+
+        original = ((original[0] + 1.0) / 2.0).detach().cpu().numpy()
+        generated = ((generated[0] + 1.0) / 2.0).detach().cpu().numpy()
+        target = ((target[0] + 1.0) / 2.0).detach().cpu().numpy()
+
+        plt.close()
+        _, axes = plt.subplots(1, 3, sharex=True, figsize=(16.5, 5.5))
+        titles = ("Original", "Starless", "Target")
+        for axis, title, image in zip(axes, titles, (original, generated, target)):
+            if self.mode == "RGB":
+                axis.imshow(np.transpose(image, (1, 2, 0)))
+            else:
+                axis.imshow(image[0], cmap="gray", vmin=0, vmax=1)
+            axis.set_title(title)
+        plt.pause(0.001)
+        self.G.train()
+
+    def train(
+        self,
+        epochs,
+        augmentation=True,
+        plot_progress=False,
+        plot_interval=50,
+        save_backups=True,
+        warm_up=False,
+    ):
+        if not self.original:
+            raise RuntimeError("Training dataset was not loaded; call load_training_dataset()")
+        if self.G is None or self.D is None:
+            self.initialize_model()
+
+        for epoch in range(epochs):
+            for iteration in range(self.iters_per_epoch):
+                original, target = self.generate_input(augmentation=augmentation)
+                original = torch.from_numpy(original.transpose(0, 3, 1, 2)).to(self.device)
+                target = torch.from_numpy(target.transpose(0, 3, 1, 2)).to(self.device)
+                original = original * 2.0 - 1.0
+                target = target * 2.0 - 1.0
+                if warm_up:
+                    target = original
+
+                if plot_progress and iteration % plot_interval == 0:
+                    self._plot_progress(original, target)
+
+                metrics = self._train_step(original, target)
+                self._update_history(metrics)
+                print(
+                    "\rEpoch: %d. Iteration %d / %d Loss %f    "
+                    % (epoch, iteration, self.iters_per_epoch, self.history["total"][-1]),
+                    end="",
+                )
+
+            print()
+            if save_backups:
+                suffix = "even" if epoch % 2 == 0 else "odd"
+                self._save_network(self.G, "./starnet_backup_G_%s.pth" % suffix)
+                self._save_network(self.D, "./starnet_backup_D_%s.pth" % suffix)
+
+        if plot_progress:
+            plt.close()
+        self.G.eval()
+        self.D.eval()
+        return self.history
+
+    def plot_history(self, last=None):
+        if not self.history:
+            raise RuntimeError("Empty training history, nothing to plot")
+
+        keys = list(self.history.keys())
+        columns = 3
+        rows = int(np.ceil(len(keys) / columns))
+        figure, axes = plt.subplots(rows, columns, sharex=True, figsize=(16, 3.5 * rows))
+        axes = np.atleast_1d(axes).reshape(-1)
+        for axis, key in zip(axes, keys):
+            values = self.history[key][-last:] if last else self.history[key]
+            axis.plot(values)
+            axis.set_title(key)
+        for axis in axes[len(keys):]:
+            axis.set_visible(False)
+        figure.tight_layout()
+        return figure, axes
+
+    @staticmethod
+    def _save_network(network, path):
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"state_dict": network.state_dict()}, path)
+
+    def save_model(self, weights_filename, history_filename=None):
+        if self.G is None or self.D is None:
+            raise RuntimeError("Model has not been initialized")
+
+        generator_path = Path(str(weights_filename) + "_G_" + self.mode + ".pth")
+        discriminator_path = Path(str(weights_filename) + "_D_" + self.mode + ".pth")
+        self._save_network(self.G, generator_path)
+        self._save_network(self.D, discriminator_path)
+
+        if history_filename:
+            history_path = Path(str(history_filename) + "_" + self.mode + ".pkl")
+            history_path.parent.mkdir(parents=True, exist_ok=True)
+            with history_path.open("wb") as handle:
+                pickle.dump(self.history, handle)
+        return generator_path, discriminator_path
+
+    def _read_image(self, input_path):
+        data = tiff.imread(input_path)
+        if data.ndim > 3:
+            layer = input("Tiff has %d layers, please enter layer to process: " % data.shape[0])
+            data = data[int(layer)]
+
+        if data.dtype == np.dtype("uint16"):
+            scale = 65535.0
+        elif data.dtype == np.dtype("uint8"):
+            scale = 255.0
+        else:
+            raise ValueError("Unknown image dtype: %s" % data.dtype)
+
+        image = (data / scale).astype(np.float32)
+        if self.mode == "Greyscale":
+            if image.ndim != 2:
+                raise ValueError("You loaded Greyscale model, but the image is RGB")
+            image = image[:, :, None]
+        else:
+            if image.ndim != 3:
+                raise ValueError("You loaded RGB model, but the image is Greyscale")
+            if image.shape[2] == 4:
+                print("Input image has 4 channels. Removing Alpha-Channel")
+                image = image[:, :, :3]
+            if image.shape[2] != 3:
+                raise ValueError("RGB input must have 3 or 4 channels")
+        return image, data.dtype, scale
+
+    @staticmethod
+    def _imwrite(path, image):
+        if hasattr(tiff, "imwrite"):
+            tiff.imwrite(path, image)
+        else:
+            tiff.imsave(path, image)
+
+    def transform(self, in_name, out_name):
+        if self.G is None:
+            raise RuntimeError("Call load_model() before transform()")
+
+        self.G.eval()
+        image, input_dtype, scale = self._read_image(in_name)
+        offset = (self.window_size - self.stride) // 2
+        h, w, _ = image.shape
+        ith = h // self.stride + 1
+        itw = w // self.stride + 1
+        dh = ith * self.stride - h
+        dw = itw * self.stride - w
+
+        image = np.concatenate((image, image[(h - dh):, :, :]), axis=0)
+        image = np.concatenate((image, image[:, (w - dw):, :]), axis=1)
+        padded_h, padded_w, _ = image.shape
+        image = np.concatenate((image, image[(padded_h - offset):, :, :]), axis=0)
+        image = np.concatenate((image[:offset, :, :], image), axis=0)
+        image = np.concatenate((image, image[:, (padded_w - offset):, :]), axis=1)
+        image = np.concatenate((image[:, :offset, :], image), axis=1)
+        image = image * 2.0 - 1.0
+        output = np.empty_like(image)
+
+        with torch.inference_mode():
+            for i in range(ith):
+                for j in range(itw):
+                    x = self.stride * i
+                    y = self.stride * j
+                    tile = image[
+                        x:x + self.window_size,
+                        y:y + self.window_size,
+                        :,
+                    ]
+                    tile = (
+                        torch.from_numpy(tile.transpose(2, 0, 1))
+                        .unsqueeze(0)
+                        .to(self.device)
+                    )
+                    result = self.G(tile)
+                    result = ((result[0] + 1.0) / 2.0).cpu().numpy().transpose(1, 2, 0)
+                    output[
+                        x + offset:x + offset + self.stride,
+                        y + offset:y + offset + self.stride,
+                        :,
+                    ] = result[
+                        offset:offset + self.stride,
+                        offset:offset + self.stride,
+                        :,
+                    ]
+                    progress = int((itw * i + j + 1) * 100 / (ith * itw))
+                    print("Transforming input image... %d%%\r" % progress, end="")
+
+        print("Transforming input image... Done!")
+        output = np.clip(output, 0.0, 1.0)
+        output = output[offset:offset + h, offset:offset + w, :]
+        if self.mode == "Greyscale":
+            output = output[:, :, 0]
+
+        output_path = Path(out_name)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        self._imwrite(output_path, (output * scale).astype(input_dtype))
+        return output_path

StarNetPyTorch/starnet_v1_pytorch_transform.py

@@ -0,0 +1,51 @@
+import argparse
+from pathlib import Path
+
+from starnet_v1_pytorch import StarNet
+
+
+def main():
+
+    input_default  = "rgb_test5.tif"
+    output_default = "rgb_test5_starless.tif"
+    weights_default = str(Path(__file__).resolve().parent / "models" / "weights")
+    mode_default = "RGB" 
+
+    parser = argparse.ArgumentParser(description="PyTorch version of the TF2 StarNet transform demo")
+    parser.add_argument("input", nargs="?", default=input_default, help="Input TIFF image")
+    parser.add_argument(
+        "output",
+        nargs="?",
+        default=output_default,
+        help="Output TIFF image",
+    )
+    parser.add_argument(
+        "--weights",
+        default=weights_default,
+        help="Weight prefix or a direct .pth file path",
+    )
+    parser.add_argument("--mode", choices=("RGB", "Greyscale"), default=mode_default)
+    parser.add_argument("--window-size", type=int, default=512)
+    parser.add_argument("--stride", type=int, default=128)
+    parser.add_argument(
+        "--device",
+        default=None,
+        help="Torch device such as cuda, cuda:0, or cpu (default: auto-detect)",
+    )
+    args = parser.parse_args()
+
+    starnet = StarNet(
+        mode=args.mode,
+        window_size=args.window_size,
+        stride=args.stride,
+        device=args.device,
+    )
+    weights_path = starnet.load_model(args.weights)
+    print("Loaded weights: %s" % weights_path)
+
+    output_path = starnet.transform(args.input, args.output)
+    print("Saved result: %s" % Path(output_path).resolve())
+
+
+if __name__ == "__main__":
+    main()

StarNetPyTorch/transform.py

@@ -0,0 +1,121 @@
+from pathlib import Path
+
+import numpy as np
+import tifffile as tiff
+import torch
+
+from model import create_generator
+
+
+WINDOW_SIZE = 256
+DEFAULT_STRIDE = 64
+
+
+def _read_tiff_image(image_name):
+    data = tiff.imread(image_name)
+    if len(data.shape) > 3:
+        layer = input("Tiff has %d layers, please enter layer to process: " % data.shape[0])
+        data = data[int(layer)]
+
+    input_dtype = data.dtype
+    if input_dtype == np.dtype("uint16"):
+        image = (data / 65535.0).astype("float32")
+    elif input_dtype == np.dtype("uint8"):
+        image = (data / 255.0).astype("float32")
+    else:
+        raise ValueError("Unknown image dtype: %s" % data.dtype)
+
+    if image.ndim != 3 or image.shape[2] < 3:
+        raise ValueError("Input image must be an RGB/RGBA TIFF image")
+
+    if image.shape[2] == 4:
+        print("Input image has 4 channels. Removing Alpha-Channel")
+        image = image[:, :, [0, 1, 2]]
+
+    return image, input_dtype
+
+
+def _default_weights_path():
+    candidates = [
+        Path.cwd() / "starnet_generator.pth",
+        Path(__file__).resolve().parent / "starnet_generator.pth",
+    ]
+    for candidate in candidates:
+        if candidate.exists():
+            return candidate
+    return candidates[0]
+
+
+def _save_tiff(path, array):
+    if hasattr(tiff, "imwrite"):
+        tiff.imwrite(path, array)
+    else:
+        tiff.imsave(path, array)
+
+
+def transform(imageName, stride=DEFAULT_STRIDE, weights_path=None, device=None):
+    weights_path = Path(weights_path) if weights_path else _default_weights_path()
+    if not weights_path.exists():
+        raise FileNotFoundError(
+            "PyTorch weights were not found at %s. Convert a TensorFlow checkpoint first "
+            "with: python convert_tf_checkpoint.py --checkpoint ../starnet/model.ckpt "
+            "--output starnet_generator.pth" % weights_path
+        )
+
+    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    print("Loading PyTorch model from %s..." % weights_path)
+    net = create_generator(weights_path=weights_path, device=device)
+    print("Done!")
+
+    print("Opening input image...")
+    image, input_dtype = _read_tiff_image(imageName)
+
+    offset = int((WINDOW_SIZE - stride) / 2)
+    h, w, _ = image.shape
+    ith = int(h / stride) + 1
+    itw = int(w / stride) + 1
+    dh = ith * stride - h
+    dw = itw * stride - w
+
+    image = np.concatenate((image, image[(h - dh):, :, :]), axis=0)
+    image = np.concatenate((image, image[:, (w - dw):, :]), axis=1)
+
+    h, w, _ = image.shape
+    image = np.concatenate((image, image[(h - offset):, :, :]), axis=0)
+    image = np.concatenate((image[:offset, :, :], image), axis=0)
+    image = np.concatenate((image, image[:, (w - offset):, :]), axis=1)
+    image = np.concatenate((image[:, :offset, :], image), axis=1)
+
+    output = np.copy(image)
+
+    with torch.inference_mode():
+        for i in range(ith):
+            for j in range(itw):
+                progress = int((itw * i + j + 1) * 100 / (ith * itw))
+                print("Transforming input image... %d%%\r" % progress, end="")
+
+                x = stride * i
+                y = stride * j
+                patch = image[x:x + WINDOW_SIZE, y:y + WINDOW_SIZE, :]
+                patch_tensor = torch.from_numpy(patch.transpose(2, 0, 1)).unsqueeze(0).to(device)
+
+                result = net(patch_tensor).squeeze(0).detach().cpu().numpy().transpose(1, 2, 0)
+                output[
+                    x + offset:x + stride + offset,
+                    y + offset:y + stride + offset,
+                    :,
+                ] = result[offset:stride + offset, offset:stride + offset, :]
+
+    print("Transforming input image... Done!")
+
+    output = np.clip(output, 0, 1)
+    output = output[offset:-(offset + dh), offset:-(offset + dw), :]
+
+    output_path = Path(str(imageName) + "_starless.tif")
+    print("Saving output image...")
+    if input_dtype == np.dtype("uint8"):
+        _save_tiff(output_path, (output * 255).astype("uint8"))
+    else:
+        _save_tiff(output_path, (output * 65535).astype("uint16"))
+    print("Done!")
+    return output_path