First commit

app.py

@@ -0,0 +1,84 @@
+import streamlit as st
+from PIL import Image
+import numpy as np
+from datasets.rg_masks import get_transforms
+from models import tiramisu
+from torchvision.transforms.functional import to_pil_image
+import torch
+from astropy.io import fits
+
+
+def load_fits(path):
+    array = fits.getdata(path).astype(np.float32)
+    array = np.expand_dims(array, 2)
+    return array
+
+def load_image(path):
+    image = Image.open(path)
+    array = np.array(image)
+    array = np.expand_dims(array[:,:,0], 2)
+
+    return array
+
+def load_weights(model, fpath, device="cuda"):
+    print("loading weights '{}'".format(fpath))
+    weights = torch.load(fpath, map_location=torch.device(device))
+    model.load_state_dict(weights['state_dict'])
+
+
+# Function to apply color overlay to the input image based on the segmentation mask
+def apply_color_overlay(input_image, segmentation_mask, alpha=0.5):
+    r = (segmentation_mask == 1).float()
+    g = (segmentation_mask == 2).float()
+    b = (segmentation_mask == 3).float()
+    overlay = torch.cat([r, g, b], dim=0)
+    overlay = to_pil_image(overlay)
+    output = Image.blend(input_image, overlay, alpha=alpha)
+    return output
+
+# Streamlit app
+def main():
+    st.title("Tiramisu for semantic segmentation of radio astronomy images")
+    st.write("Upload an image and see the segmentation result!")
+
+    uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png", "fits"])
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = tiramisu.FCDenseNet67(n_classes=4).to(device)
+    load_weights(model, "weights/real.pth")
+    model.eval()
+
+    st.markdown(
+        """
+        Category Legend:
+        - :blue[Extended]
+        - :green[Compact]
+        - :red[Spurious]
+        """
+        )
+    if uploaded_image is not None:
+        # Load the uploaded image
+        if uploaded_image.name.endswith(".fits"):
+            input_array = load_fits(uploaded_image)
+        else:
+            input_array = load_image(uploaded_image)
+
+        input_array = input_array.transpose(2,0,1)
+        transforms = get_transforms(input_array.shape[1])
+        image = transforms(input_array)
+        image = image.to(device)
+
+        with torch.no_grad():
+            output = model(image)
+        preds = output.argmax(1)
+
+        pil_image = to_pil_image(image[0])
+        # Apply color overlay to the input image
+        segmented_image = apply_color_overlay(pil_image, preds)
+
+        # Display the input image and the segmented output
+        st.image([pil_image, segmented_image], caption=["Input Image", "Segmented Output"], width=300)
+
+if __name__ == "__main__":
+    main()

datasets/rg_masks.py

@@ -0,0 +1,326 @@
+import json
+import math
+import random
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.data
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+from astropy.io import fits
+from astropy.io.fits.verify import VerifyWarning
+from einops import rearrange
+from torch.utils.data import Dataset
+from torchvision.transforms.functional import to_pil_image
+from torchvision.utils import make_grid, save_image
+
+warnings.simplefilter('ignore', category=VerifyWarning)
+import warnings
+
+import numpy as np
+import torch
+from astropy.stats import sigma_clip
+from astropy.visualization import ZScaleInterval
+from torch.utils.data import DataLoader
+
+warnings.simplefilter('ignore', category=VerifyWarning)
+
+
+CLASSES = ['background', 'spurious', 'compact', 'extended']
+COLORS = [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]
+
+
+def get_transforms(img_size):
+    return  T.Compose([
+                RemoveNaNs(),
+                ZScale(),
+                SigmaClip(),
+                ToTensor(),
+                torch.nn.Tanh(),
+                MinMaxNormalize(),
+                Unsqueeze(),
+                T.Resize((img_size, img_size)),
+                RepeatChannels((3))
+            ])
+
+class RemoveNaNs(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        img[np.isnan(img)] = 0
+        return img
+
+
+class ZScale(object):
+    def __init__(self, contrast=0.15):
+        self.contrast = contrast
+
+    def __call__(self, img):
+        interval = ZScaleInterval(contrast=self.contrast)
+        min, max = interval.get_limits(img)
+
+        img = (img - min) / (max - min)
+        return img
+
+
+class SigmaClip(object):
+    def __init__(self, sigma=3, masked=True):
+        self.sigma = sigma
+        self.masked = masked
+
+    def __call__(self, img):
+        img = sigma_clip(img, sigma=self.sigma, masked=self.masked)
+        return img
+
+
+class MinMaxNormalize(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        img = (img - img.min()) / (img.max() - img.min())
+        return img
+
+
+class ToTensor(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        return torch.tensor(img, dtype=torch.float32)
+
+class RepeatChannels(object):
+    def __init__(self, ch):
+        self.ch = ch
+
+    def __call__(self, img):
+        return img.repeat(1, self.ch, 1, 1)
+
+class FromNumpy(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        return torch.from_numpy(img.astype(np.float32)).type(torch.float32)
+
+class Unsqueeze(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        return img.unsqueeze(0)
+
+
+def mask_to_rgb(mask):
+    rgb_mask = torch.zeros_like(mask, device=mask.device).repeat(1, 3, 1, 1)
+    for i, c in enumerate(COLORS):
+        color_mask = torch.tensor(c, device=mask.device).unsqueeze(
+            1).unsqueeze(2) * (mask == i)
+        rgb_mask += color_mask
+    return rgb_mask
+
+def get_data_loader(dataset, batch_size, split="train"):
+    batch_size = batch_size
+    workers = min(8, batch_size)
+    is_train = split == "train"
+    return DataLoader(dataset, shuffle=is_train, batch_size=batch_size,
+                      num_workers=workers, persistent_workers=True,
+                      drop_last=is_train
+                      )
+
+def rgb_to_tensor(mask):
+    r,g,b = mask
+    r *= 1
+    g *= 2
+    b *= 3
+    mask, _ = torch.max(torch.stack([r,g,b]), dim=0, keepdim=True)
+    return mask
+
+
+def rand_horizontal_flip(img, mask):
+    if random.random() < 0.5:
+        img = TF.hflip(img)
+        mask = TF.hflip(mask)
+    return img, mask
+
+
+class RGDataset(Dataset):
+    def __init__(self, data_dir, img_paths, img_size=128):
+        super().__init__()
+        data_dir = Path(data_dir)
+        with open(img_paths) as f:
+            self.img_paths = f.read().splitlines()
+        self.img_paths = [data_dir / p for p in self.img_paths]
+
+        self.transforms = T.Compose([
+            RemoveNaNs(),
+            ZScale(),
+            SigmaClip(),
+            ToTensor(),
+            torch.nn.Tanh(),
+            MinMaxNormalize(),
+            # T.Resize((img_size),
+            #          interpolation=T.InterpolationMode.NEAREST),
+            Unsqueeze(),
+            T.Resize((img_size, img_size)),
+            
+            RepeatChannels((3))
+        ])
+        self.img_size = img_size
+
+        self.mask_transforms = T.Compose([
+            FromNumpy(),
+            Unsqueeze(),
+            T.Resize((img_size, img_size),
+                     interpolation=T.InterpolationMode.NEAREST),
+        ])
+
+    def get_mask(self, img_path, type):
+        assert type in ["real", "synthetic"], f"Type {type} not supported"
+        if type == "real":
+            ann_path = str(img_path).replace(
+                'imgs', 'masks').replace('.fits', '.json')
+            ann_dir = Path(ann_path).parent
+            ann_path = ann_dir / f'mask_{ann_path.split("/")[-1]}'
+            with open(ann_path) as j:
+                mask_info = json.load(j)
+
+            masks = []
+
+            for obj in mask_info['objs']:
+                seg_path = ann_dir / obj['mask']
+
+                mask = fits.getdata(seg_path)
+
+                mask = self.mask_transforms(mask.astype(np.float32))
+                masks.append(mask)
+            mask, _ = torch.max(torch.stack(masks), dim=0)
+
+        elif type == "synthetic":
+            mask_path = str(img_path).replace("gen_fits", "cond_fits")
+            mask = fits.getdata(mask_path)
+            mask = self.mask_transforms(mask)
+            mask = mask.squeeze()
+            if mask.shape[0] == 3:
+                mask = rgb_to_tensor(mask)
+        return mask
+
+
+    def __len__(self):
+        return len(self.img_paths)
+
+    def __getitem__(self, idx):
+        image_path = self.img_paths[idx]
+        img = fits.getdata(image_path)
+        img = self.transforms(img)
+        
+        if "synthetic" in str(image_path):
+            mask = self.get_mask(image_path, type='synthetic')
+        else:
+            mask = self.get_mask(image_path, type='real')
+
+        # ann_path = str(image_path).replace(
+        #     'imgs', 'masks').replace('.fits', '.json')
+        # ann_dir = Path(ann_path).parent
+        # ann_path = ann_dir / f'mask_{ann_path.split("/")[-1]}'
+        # with open(ann_path) as j:
+        #     mask_info = json.load(j)
+
+
+        # masks = []
+
+        # for obj in mask_info['objs']:
+        #     seg_path = ann_dir / obj['mask']
+
+        #     mask = fits.getdata(seg_path)
+
+        #     mask = self.mask_transforms(mask.astype(np.float32))
+            # masks.append(mask)
+
+        # if 'bkg' in str(image_path):
+        #     mask = torch.zeros_like(img)
+        #     masks.append(mask)
+
+        # mask, _ = torch.max(torch.stack(masks), dim=0)
+        mask = mask.long()
+        return img.squeeze(), mask.squeeze()
+
+
+class SyntheticRGDataset(Dataset):
+    def __init__(self, data_dir, img_paths, img_size=128):
+        super().__init__()
+        data_dir = Path(data_dir)
+        with open(img_paths) as f:
+            self.img_paths = f.read().splitlines()
+        self.img_paths = [data_dir / p for p in self.img_paths]
+
+
+
+        self.transforms = T.Compose([
+            RemoveNaNs(),
+            ZScale(),
+            SigmaClip(),
+            ToTensor(),
+            torch.nn.Tanh(),
+            MinMaxNormalize(),
+            # T.Resize((img_size),
+            #          interpolation=T.InterpolationMode.NEAREST),
+            Unsqueeze(),
+            T.Resize((img_size, img_size)),
+            
+            RepeatChannels((3))
+        ])
+        self.img_size = img_size
+
+        self.mask_transforms = T.Compose([
+            FromNumpy(),
+            Unsqueeze(),
+            T.Resize((img_size, img_size),
+                     interpolation=T.InterpolationMode.NEAREST),
+        ])
+
+    def __len__(self):
+        return len(self.img_paths)
+
+    def __getitem__(self, idx):
+        image_path = self.img_paths[idx]
+        img = fits.getdata(image_path)
+        img = self.transforms(img)
+        img = img.squeeze()
+
+        mask_path = str(image_path).replace("gen_fits", "cond_fits")
+        mask = fits.getdata(mask_path)
+        mask = self.mask_transforms(mask)
+
+        img, mask = rand_horizontal_flip(img, mask)
+
+        mask = mask.squeeze().long()
+        return img, mask
+
+
+if __name__ == '__main__':
+    rgtrain = SyntheticRGDataset('data/rg-dataset/data',
+                        'data/rg-dataset/val_w_bg.txt')
+    batch = next(iter(rgtrain))
+    image, mask, masked_image = batch
+    to_pil_image(image).save('image.png')
+    rgb_mask = mask_to_rgb(mask)[0]
+    to_pil_image(rgb_mask).save('mask.png')
+    to_pil_image(masked_image[0]).save('masked.png')
+
+    bs = 256
+
+    loader = torch.utils.data.DataLoader(
+        rgtrain, batch_size=bs, shuffle=False, num_workers=16)
+    for i, batch in enumerate(loader):
+        image, mask, masked_image = batch
+        rgb_mask = mask_to_rgb(mask)
+        nrow = int(math.sqrt(bs))
+        # nrow = bs // 2
+        grid = make_grid(rgb_mask, nrow=nrow, padding=0)
+        save_image(grid, f'mask_{nrow}x{nrow}.png')
+        break

models/layers.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+
+
+class DenseLayer(nn.Sequential):
+    def __init__(self, in_channels, growth_rate):
+        super().__init__()
+        self.add_module('norm', nn.BatchNorm2d(in_channels))
+        self.add_module('relu', nn.ReLU(True))
+        self.add_module('conv', nn.Conv2d(in_channels, growth_rate, kernel_size=3,
+                                          stride=1, padding=1, bias=True))
+        self.add_module('drop', nn.Dropout2d(0.2))
+
+    def forward(self, x):
+        return super().forward(x)
+
+
+class DenseBlock(nn.Module):
+    def __init__(self, in_channels, growth_rate, n_layers, upsample=False):
+        super().__init__()
+        self.upsample = upsample
+        self.layers = nn.ModuleList([DenseLayer(
+            in_channels + i*growth_rate, growth_rate)
+            for i in range(n_layers)])
+
+    def forward(self, x):
+        if self.upsample:
+            new_features = []
+            # we pass all previous activations into each dense layer normally
+            # But we only store each dense layer's output in the new_features array
+            for layer in self.layers:
+                out = layer(x)
+                x = torch.cat([x, out], 1)
+                new_features.append(out)
+            return torch.cat(new_features, 1)
+        else:
+            for layer in self.layers:
+                out = layer(x)
+                x = torch.cat([x, out], 1)  # 1 = channel axis
+            return x
+
+
+class TransitionDown(nn.Sequential):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.add_module('norm', nn.BatchNorm2d(num_features=in_channels))
+        self.add_module('relu', nn.ReLU(inplace=True))
+        self.add_module('conv', nn.Conv2d(in_channels, in_channels,
+                                          kernel_size=1, stride=1,
+                                          padding=0, bias=True))
+        self.add_module('drop', nn.Dropout2d(0.2))
+        self.add_module('maxpool', nn.MaxPool2d(2))
+
+    def forward(self, x):
+        return super().forward(x)
+
+
+class TransitionUp(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.convTrans = nn.ConvTranspose2d(
+            in_channels=in_channels, out_channels=out_channels,
+            kernel_size=3, stride=2, padding=0, bias=True)
+
+    def forward(self, x, skip):
+        out = self.convTrans(x)
+        out = center_crop(out, skip.size(2), skip.size(3))
+        out = torch.cat([out, skip], 1)
+        return out
+
+
+class Bottleneck(nn.Sequential):
+    def __init__(self, in_channels, growth_rate, n_layers):
+        super().__init__()
+        self.add_module('bottleneck', DenseBlock(
+            in_channels, growth_rate, n_layers, upsample=True))
+
+    def forward(self, x):
+        return super().forward(x)
+
+
+def center_crop(layer, max_height, max_width):
+    _, _, h, w = layer.size()
+    xy1 = (w - max_width) // 2
+    xy2 = (h - max_height) // 2
+    return layer[:, :, xy2:(xy2 + max_height), xy1:(xy1 + max_width)]

models/tiramisu.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+
+from .layers import *
+
+
+class FCDenseNet(nn.Module):
+    def __init__(self, in_channels=3, down_blocks=(5, 5, 5, 5, 5),
+                 up_blocks=(5, 5, 5, 5, 5), bottleneck_layers=5,
+                 growth_rate=16, out_chans_first_conv=48, n_classes=12):
+        super().__init__()
+        self.down_blocks = down_blocks
+        self.up_blocks = up_blocks
+        cur_channels_count = 0
+        skip_connection_channel_counts = []
+
+        ## First Convolution ##
+
+        self.add_module('firstconv', nn.Conv2d(in_channels=in_channels,
+                                               out_channels=out_chans_first_conv, kernel_size=3,
+                                               stride=1, padding=1, bias=True))
+        cur_channels_count = out_chans_first_conv
+
+        #####################
+        # Downsampling path #
+        #####################
+
+        self.denseBlocksDown = nn.ModuleList([])
+        self.transDownBlocks = nn.ModuleList([])
+        for i in range(len(down_blocks)):
+            self.denseBlocksDown.append(
+                DenseBlock(cur_channels_count, growth_rate, down_blocks[i]))
+            cur_channels_count += (growth_rate*down_blocks[i])
+            skip_connection_channel_counts.insert(0, cur_channels_count)
+            self.transDownBlocks.append(TransitionDown(cur_channels_count))
+
+        #####################
+        #     Bottleneck    #
+        #####################
+
+        self.add_module('bottleneck', Bottleneck(cur_channels_count,
+                                                 growth_rate, bottleneck_layers))
+        prev_block_channels = growth_rate*bottleneck_layers
+        cur_channels_count += prev_block_channels
+
+        #######################
+        #   Upsampling path   #
+        #######################
+
+        self.transUpBlocks = nn.ModuleList([])
+        self.denseBlocksUp = nn.ModuleList([])
+        for i in range(len(up_blocks)-1):
+            self.transUpBlocks.append(TransitionUp(
+                prev_block_channels, prev_block_channels))
+            cur_channels_count = prev_block_channels + \
+                skip_connection_channel_counts[i]
+
+            self.denseBlocksUp.append(DenseBlock(
+                cur_channels_count, growth_rate, up_blocks[i],
+                upsample=True))
+            prev_block_channels = growth_rate*up_blocks[i]
+            cur_channels_count += prev_block_channels
+
+        ## Final DenseBlock ##
+
+        self.transUpBlocks.append(TransitionUp(
+            prev_block_channels, prev_block_channels))
+        cur_channels_count = prev_block_channels + \
+            skip_connection_channel_counts[-1]
+
+        self.denseBlocksUp.append(DenseBlock(
+            cur_channels_count, growth_rate, up_blocks[-1],
+            upsample=False))
+        cur_channels_count += growth_rate*up_blocks[-1]
+
+        ## Softmax ##
+
+        self.finalConv = nn.Conv2d(in_channels=cur_channels_count,
+                                   out_channels=n_classes, kernel_size=1, stride=1,
+                                   padding=0, bias=True)
+        self.softmax = nn.LogSoftmax(dim=1)
+
+    def forward(self, x):
+        out = self.firstconv(x)
+
+        skip_connections = []
+        for i in range(len(self.down_blocks)):
+            out = self.denseBlocksDown[i](out)
+            skip_connections.append(out)
+            out = self.transDownBlocks[i](out)
+
+        out = self.bottleneck(out)
+        for i in range(len(self.up_blocks)):
+            skip = skip_connections.pop()
+            out = self.transUpBlocks[i](out, skip)
+            out = self.denseBlocksUp[i](out)
+
+        out = self.finalConv(out)
+        out = self.softmax(out)
+        return out
+
+
+def FCDenseNet57(n_classes):
+    return FCDenseNet(
+        in_channels=3, down_blocks=(4, 4, 4, 4, 4),
+        up_blocks=(4, 4, 4, 4, 4), bottleneck_layers=4,
+        growth_rate=12, out_chans_first_conv=48, n_classes=n_classes)
+
+
+def FCDenseNet67(n_classes):
+    return FCDenseNet(
+        in_channels=3, down_blocks=(5, 5, 5, 5, 5),
+        up_blocks=(5, 5, 5, 5, 5), bottleneck_layers=5,
+        growth_rate=16, out_chans_first_conv=48, n_classes=n_classes)
+
+
+def FCDenseNet103(n_classes):
+    return FCDenseNet(
+        in_channels=3, down_blocks=(4, 5, 7, 10, 12),
+        up_blocks=(12, 10, 7, 5, 4), bottleneck_layers=15,
+        growth_rate=16, out_chans_first_conv=48, n_classes=n_classes)