models/model.py

import random
from statistics import mean
from typing import List, Tuple

import torch as th
import pytorch_lightning as pl
from jaxtyping import Float, Int
import numpy as np
from torch_geometric.nn.conv import GATv2Conv

from models.SAP.dpsr import DPSR
from models.SAP.model import PSR2Mesh

# Constants

th.manual_seed(0)
np.random.seed(0)

BATCH_SIZE = 1 # BS

IN_DIM = 1 
OUT_DIM = 1
LATENT_DIM = 32

DROPOUT_PROB = 0.1
GRID_SIZE = 128

def generate_grid_edge_list(gs: int = 128):
    grid_edge_list = []

    for k in range(gs):
        for j in range(gs):
            for i in range(gs):
                current_idx = i + gs*j + k*gs*gs
                if (i - 1) >= 0:
                    grid_edge_list.append([current_idx, i-1 + gs*j + k*gs*gs])
                if (i + 1) < gs:
                    grid_edge_list.append([current_idx, i+1 + gs*j + k*gs*gs])
                if (j - 1) >= 0:
                    grid_edge_list.append([current_idx, i + gs*(j-1) + k*gs*gs])
                if (j + 1) < gs:
                    grid_edge_list.append([current_idx, i + gs*(j+1) + k*gs*gs])
                if (k - 1) >= 0:
                    grid_edge_list.append([current_idx, i + gs*j + (k-1)*gs*gs])
                if (k + 1) < gs:
                    grid_edge_list.append([current_idx, i + gs*j + (k+1)*gs*gs])
    return grid_edge_list

GRID_EDGE_LIST = generate_grid_edge_list(GRID_SIZE)
GRID_EDGE_LIST = th.tensor(GRID_EDGE_LIST, dtype=th.int)
GRID_EDGE_LIST = GRID_EDGE_LIST.T
# GRID_EDGE_LIST = GRID_EDGE_LIST.to(th.device("cuda"))
GRID_EDGE_LIST.requires_grad = False # Do not forget to delete it if train


class FormOptimizer(th.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        
        layers = []
        
        self.gconv1 = GATv2Conv(in_channels=IN_DIM, out_channels=LATENT_DIM, heads=1, dropout=DROPOUT_PROB)
        self.gconv2 = GATv2Conv(in_channels=LATENT_DIM, out_channels=LATENT_DIM, heads=1, dropout=DROPOUT_PROB)
        
        self.actv = th.nn.Sigmoid()
        self.head = th.nn.Linear(in_features=LATENT_DIM, out_features=OUT_DIM)

    def forward(self, 
                field: Float[th.Tensor, "GS GS GS"]) -> Float[th.Tensor, "GS GS GS"]:
        """
         Args:
        field (Tensor [GS, GS, GS]): vertices and normals tensor.
        """
        vertex_features = field.clone()
        vertex_features = vertex_features.reshape(GRID_SIZE*GRID_SIZE*GRID_SIZE, IN_DIM)
        
        vertex_features = self.gconv1(x=vertex_features, edge_index=GRID_EDGE_LIST) 
        vertex_features = self.gconv2(x=vertex_features, edge_index=GRID_EDGE_LIST) 
        field_delta = self.head(self.actv(vertex_features))
        
        field_delta = field_delta.reshape(BATCH_SIZE, GRID_SIZE, GRID_SIZE, GRID_SIZE)
        field_delta += field # field_delta carries the gradient
        field_delta = th.clamp(field_delta, min=-0.5, max=0.5)
        
        return field_delta

class Model(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.form_optimizer = FormOptimizer()
        
        self.dpsr =  DPSR([GRID_SIZE, GRID_SIZE, GRID_SIZE], sig=0.0)
        self.field2mesh = PSR2Mesh().apply

        self.metric = th.nn.MSELoss()
            
        self.val_losses = []
        self.train_losses = []

    def log_h5(self, points, normals):
        dset = self.log_points_file.create_dataset(
                                    name=str(self.h5_frame),
                                    shape=points.shape,
                                    dtype=np.float16, 
                                    compression="gzip")
        dset[:] = points
        dset = self.log_normals_file.create_dataset(
                                     name=str(self.h5_frame),
                                     shape=normals.shape,
                                     dtype=np.float16, 
                                     compression="gzip")
        dset[:] = normals
        self.h5_frame += 1
    
    def forward(self, 
                v: Float[th.Tensor, "BS N 3"],
                n: Float[th.Tensor, "BS N 3"]) -> Tuple[Float[th.Tensor, "BS N 3"], # v - vertices
                                                        Int[th.Tensor, "2 E"], # f - faces
                                                        Float[th.Tensor, "BS N 3"], # n - vertices normals
                                                        Float[th.Tensor, "BS GR GR GR"]]: # field: 
        field = self.dpsr(v, n)
        field = self.form_optimizer(field)
        v, f, n = self.field2mesh(field)
        return v, f, n, field

    def training_step(self, batch, batch_idx) -> Float[th.Tensor, "1"]:
        vertices, vertices_normals, vertices_gt, vertices_normals_gt, field_gt, adj = batch
            
        mask = th.rand((vertices.shape[1], ), device=th.device("cuda")) < (random.random() / 2.0 + 0.5)
        vertices = vertices[:, mask]
        vertices_normals = vertices_normals[:, mask]
        
        vr, fr, nr, field_r = model(vertices, vertices_normals)
        
        loss = self.metric(field_r, field_gt)
        train_per_step_loss = loss.item()
        self.train_losses.append(train_per_step_loss)
        
        return loss
    
    def on_train_epoch_end(self):
        mean_train_per_epoch_loss = mean(self.train_losses)
        self.log("mean_train_per_epoch_loss", mean_train_per_epoch_loss, on_step=False, on_epoch=True)
        self.train_losses = []
    
    def validation_step(self, batch, batch_idx):
        vertices, vertices_normals, vertices_gt, vertices_normals_gt, field_gt, adj = batch
            
        vr, fr, nr, field_r = model(vertices, vertices_normals)
        
        loss = self.metric(field_r, field_gt)
        val_per_step_loss = loss.item()
        self.val_losses.append(val_per_step_loss)
        return loss
    
    def on_validation_epoch_end(self):
        mean_val_per_epoch_loss = mean(self.val_losses)
        self.log("mean_val_per_epoch_loss", mean_val_per_epoch_loss, on_step=False, on_epoch=True)
        self.val_losses = []

    def configure_optimizers(self):
        optimizer = th.optim.Adam(self.parameters(), lr=LR)
        scheduler = th.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
        
        return {
                "optimizer": optimizer,
                "lr_scheduler": {
                                 "scheduler": scheduler, 
                                 "monitor": "mean_val_per_epoch_loss",
                                 "interval": "epoch",
                                 "frequency": 1,
                                 # If set to `True`, will enforce that the value specified 'monitor'
                                 # is available when the scheduler is updated, thus stopping
                                 # training if not found. If set to `False`, it will only produce a warning
                                 "strict": True,
                                 # If using the `LearningRateMonitor` callback to monitor the
                                 # learning rate progress, this keyword can be used to specify
                                 # a custom logged name
                                 "name": None,
                                }
                }