shape_complexity.py

import os

# from zlib import compress
from bz2 import compress

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import torch
import torch.nn.functional as F
from kornia.morphology import closing
from PIL import Image
from torch import Tensor, nn
from torch.optim import Adam
from torch.utils.data import DataLoader, RandomSampler
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms
from torchvision.utils import make_grid, save_image

device = torch.device("cuda")
matplotlib.use("Agg")

dx = [+1, 0, -1, 0]
dy = [0, +1, 0, -1]


# perform depth first search for each candidate/unlabeled region
# reference: https://stackoverflow.com/questions/14465297/connected-component-labeling-implementation
def dfs(mask: npt.NDArray, x: int, y: int, labels: npt.NDArray, current_label: int):
    n_rows, n_cols = mask.shape
    if x < 0 or x == n_rows:
        return
    if y < 0 or y == n_cols:
        return
    if labels[x][y] or not mask[x][y]:
        return  # already labeled or not marked with 1 in image

    # mark the current cell
    labels[x][y] = current_label

    # recursively mark the neighbors
    for direction in range(4):
        dfs(mask, x + dx[direction], y + dy[direction], labels, current_label)


def find_components(mask: npt.NDArray):
    label = 0

    n_rows, n_cols = mask.shape
    labels = np.zeros(mask.shape, dtype=np.int8)

    for i in range(n_rows):
        for j in range(n_cols):
            if not labels[i][j] and mask[i][j]:
                label += 1
                dfs(mask, i, j, labels, label)

    return labels


# https://stackoverflow.com/questions/31400769/bounding-box-of-numpy-array
def bbox(img):
    max_x, max_y = img.shape
    rows = np.any(img, axis=1)
    cols = np.any(img, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]

    rmin = rmin - 1 if rmin > 0 else rmin
    cmin = cmin - 1 if cmin > 0 else cmin
    rmax = rmax + 1 if rmax < max_x else rmax
    cmax = cmax + 1 if cmax < max_y else cmax

    return rmin, rmax, cmin, cmax


def extract_single_masks(labels: npt.NDArray):
    masks = []
    for l in range(labels.max() + 1):
        mask = (labels == l).astype(np.int8)
        rmin, rmax, cmin, cmax = bbox(mask)
        masks.append(mask[rmin : rmax + 1, cmin : cmax + 1])

    return masks


class VAE(nn.Module):
    """
    https://github.com/pytorch/examples/blob/main/vae/main.py
    """

    def __init__(self, bottleneck=2, image_dim=4096):
        super(VAE, self).__init__()

        self.bottleneck = bottleneck
        self.image_dim = image_dim

        self.prelim_encode = nn.Sequential(
            nn.Flatten(), nn.Linear(image_dim, 400), nn.ReLU()
        )
        self.encode_mu = nn.Sequential(nn.Linear(400, bottleneck))
        self.encode_logvar = nn.Sequential(nn.Linear(400, bottleneck))

        self.decode = nn.Sequential(
            nn.Linear(bottleneck, 400),
            nn.ReLU(),
            nn.Linear(400, image_dim),
            nn.Sigmoid(),
        )

    def encode(self, x):
        # h1 = F.relu(self.encode(x))
        # return self.encode_mu(h1), self.encode_logvar(h1)
        x = self.prelim_encode(x)
        return self.encode_mu(x), self.encode_logvar(x)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

    # Reconstruction + KL divergence losses summed over all elements and batch
    def loss(self, recon_x, x, mu, logvar):
        BCE = F.binary_cross_entropy(recon_x, x.view(-1, 4096), reduction="sum")

        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

        return BCE + KLD


class CONVVAE(nn.Module):
    def __init__(
        self,
        bottleneck=2,
    ):
        super(CONVVAE, self).__init__()

        self.bottleneck = bottleneck
        self.feature_dim = 6 * 6 * 64

        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),  # -> 30x30x16
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),  # -> 14x14x32
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(32, 64, 3),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),  # -> 6x6x64
        )
        # self.conv4 = nn.Sequential(
        #     nn.Conv2d(32, self.bottleneck, 5),
        #     nn.ReLU(),
        #     nn.MaxPool2d((2, 2), return_indices=True),  # -> 1x1xbottleneck
        # )

        self.encode_mu = nn.Sequential(
            nn.Flatten(),
            nn.Linear(self.feature_dim, self.bottleneck),
        )
        self.encode_logvar = nn.Sequential(
            nn.Flatten(), nn.Linear(self.feature_dim, self.bottleneck)
        )

        self.decode_linear = nn.Linear(self.bottleneck, self.feature_dim)

        # self.decode4 = nn.Sequential(
        #     nn.ConvTranspose2d(self.bottleneck, 32, 5),
        #     nn.ReLU(),
        # )
        self.decode3 = nn.Sequential(
            nn.ConvTranspose2d(64, 64, 2, stride=2),  # -> 12x12x64
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 3),  # -> 14x14x32
            nn.ReLU(),
        )
        self.decode2 = nn.Sequential(
            nn.ConvTranspose2d(32, 32, 2, stride=2),  # -> 28x28x32
            nn.ReLU(),
            nn.ConvTranspose2d(32, 16, 3),  # -> 30x30x16
            nn.ReLU(),
        )
        self.decode1 = nn.Sequential(
            nn.ConvTranspose2d(16, 16, 2, stride=2),  # -> 60x60x16
            nn.ReLU(),