crisgi 0.1.0__py3-none-any.whl

crisgi/__init__.py

@@ -0,0 +1,4 @@
+from .crisgi import CRISGITime
+from ._version import __version__
+
+__all__ = ['CRISGITime', '__version__']

crisgi/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

crisgi/cnn/CNNModel.py

@@ -0,0 +1,64 @@
+import sys
+import torch
+from crisgi.cnn.autoEncoder import AE
+import crisgi.cnn.autoEncoder
+from crisgi.cnn.mlp import MLP
+import crisgi.cnn.mlp
+from crisgi.cnn.train import train as train_loop
+
+sys.modules['autoEncoder'] = crisgi.cnn.autoEncoder
+sys.modules['mlp'] = crisgi.cnn.mlp
+
+
+class CNNModel:
+    def __init__(self, device, ae_path=None, mlp_path=None):
+        self.device = device
+        self.ae = AE().to(device)
+        self.mlp = MLP().to(device)
+        self.ae_loss_fn = torch.nn.MSELoss().to(device)
+        self.ce_loss_fn = torch.nn.CrossEntropyLoss().to(device)
+        self.optimizer = torch.optim.Adam(
+            list(self.ae.parameters()) + list(self.mlp.parameters()), 
+            lr=0.001, weight_decay=1e-4
+        )
+
+        if ae_path and mlp_path:
+            print(f"Loading pretrained models:\n - AE: {ae_path}\n - MLP: {mlp_path}")
+            self.ae = torch.load(ae_path, map_location=device, weights_only=False)
+            self.mlp = torch.load(mlp_path, map_location=device, weights_only=False)
+
+
+    def train(self, train_loader,epochs=10):
+        final_metrics = None
+        for epoch in range(epochs):
+            print(f"\nEpoch {epoch + 1}/{epochs}")
+            final_metrics = train_loop(
+                self.ae,
+                self.mlp,
+                train_loader,
+                self.ae_loss_fn,
+                self.ce_loss_fn,
+                self.optimizer,
+                self.device
+            )
+        return final_metrics
+
+    def predict(self, data_loader):
+        self.ae.eval()
+        self.mlp.eval()
+        all_predictions = []
+
+        with torch.no_grad():
+            for x, _ in data_loader:  
+                x = x.to(self.device)
+                en, _ = self.ae(x)
+                out = self.mlp(en)
+                predicted = out.argmax(dim=1)
+                all_predictions.extend(predicted.cpu().numpy())
+
+        return all_predictions
+    
+    def save(self, ae_path, mlp_path):
+        torch.save(self.ae.state_dict(), ae_path)
+        torch.save(self.mlp.state_dict(), mlp_path)
+        print(f"Models saved to:\n - {ae_path}\n - {mlp_path}")

crisgi/cnn/autoEncoder.py

@@ -0,0 +1,36 @@
+from torch import nn
+
+
+class AE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(4, 16, 3, padding=1),
+            nn.BatchNorm2d(16),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.MaxPool2d(2, 2),
+            nn.Conv2d(16, 32, 3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.MaxPool2d(2, 2),
+            nn.Conv2d(32, 64, 3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.MaxPool2d(2, 2)
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(16, 4, 3, stride=2, padding=1, output_padding=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        en = self.encoder(x)
+        de = self.decoder(en)
+        return en, de

crisgi/cnn/evalution_metrics.py

@@ -0,0 +1,46 @@
+from sklearn.metrics import auc, brier_score_loss, cohen_kappa_score, confusion_matrix, f1_score, roc_curve, roc_auc_score, precision_recall_curve, average_precision_score
+
+
+def calculate_pred_metric(label, pred):
+    # label: ground truth
+    # pred: prediction
+
+    # calculate ROC AUC
+    roc_auc = roc_auc_score(label, pred)
+
+    # confusion matrix
+    tn, fp, fn, tp = confusion_matrix(label, pred).ravel()
+
+    # other metric
+    accuracy = (tp + tn) / (tp + fp + fn + tn)
+    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
+    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
+    sensitivity = recall
+    specificity = tn / (tn + fp)
+    ppv = precision
+    npv = tn / (tn + fn) if (tn + fn) > 0 else 0
+    f1 = f1_score(label, pred)
+    kappa = cohen_kappa_score(label, pred)
+    brier = brier_score_loss(label, pred)
+
+    # calculate PR curve and its AUC
+    precision_vals, recall_vals, _ = precision_recall_curve(label, pred)
+    avg_pr_auc = auc(recall_vals, precision_vals)
+
+    # store the result into dictionary
+    res = {
+        'AUROC': roc_auc,
+        'AUPRC': avg_pr_auc,
+        'Accuracy': accuracy,
+        'Precision': precision,
+        'Recall': recall,
+        'Sensitivity': sensitivity,
+        'Specificity': specificity,
+        'PPV': ppv,
+        'NPV': npv,
+        'F1_Score': f1,
+        'Kappa': kappa,
+        'Brier_Score': brier
+    }
+
+    return res

crisgi/cnn/mlp.py

@@ -0,0 +1,16 @@
+from torch import nn
+
+
+class MLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(64 * 28 * 28, 64),
+            nn.ReLU(),
+            nn.Linear(64, 2),
+        )
+
+    def forward(self, x):
+        flatt = x.view(x.size(0), -1)
+        res = self.mlp(flatt)
+        return res

crisgi/cnn/train.py

@@ -0,0 +1,42 @@
+from crisgi.cnn.evalution_metrics import calculate_pred_metric
+
+
+def train(ae, mlp, train_loader, AE_loss_function, CE_loss_function, optimizer, device):
+    ae.train()
+    mlp.train()
+    total_loss = 0
+    correct = 0
+    size = len(train_loader.dataset)
+    all_predictions = []
+    all_labels = []
+
+    for x, y in train_loader:
+        x = x.to(device)
+        y = y.to(device)
+        optimizer.zero_grad()
+        en, de = ae(x)
+        classification_res = mlp(en)
+
+        AE_loss = AE_loss_function(de, x)
+        CE_loss = CE_loss_function(classification_res, y)
+        loss = AE_loss + CE_loss
+
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item()
+        predicted = classification_res.argmax(1)
+        correct += (predicted == y).sum().item()
+
+        all_labels.extend(y.cpu().numpy())
+        all_predictions.extend(predicted.cpu().numpy())
+
+    avg_loss = total_loss / size
+    accuracy = 100 * correct / size
+
+    print(f"Total Train Loss: {avg_loss}")
+    print(f"Train Accuracy: {accuracy}%")
+
+    metrics = calculate_pred_metric(all_labels, all_predictions)
+
+    return metrics