src/images/Diffusion/diffusion_model_classifier.py

import argparse
import logging
import os

import pandas as pd
import pytorch_lightning as pl
import timm
import torch
import torchvision.transforms as transforms
from data_split import *
from dataloader import *
from PIL import Image
from pytorch_lightning.callbacks import (
    EarlyStopping,
    ModelCheckpoint,
)
from sklearn.metrics import roc_auc_score
from torchmetrics import (
    Accuracy,
    Recall,
)
from utils_sampling import *

logging.basicConfig(
    filename="training.log", filemode="w", level=logging.INFO, force=True
)


class ImageClassifier(pl.LightningModule):
    def __init__(self, lmd=0):
        super().__init__()
        self.model = timm.create_model(
            "resnet50", pretrained=True, num_classes=1
        )
        self.accuracy = Accuracy(task="binary", threshold=0.5)
        self.recall = Recall(task="binary", threshold=0.5)
        self.validation_outputs = []
        self.lmd = lmd

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch):
        images, labels, _ = batch
        outputs = self.forward(images).squeeze()

        print(f"Shape of outputs (training): {outputs.shape}")
        print(f"Shape of labels (training): {labels.shape}")

        loss = F.binary_cross_entropy_with_logits(outputs, labels.float())
        logging.info(f"Training Step - ERM loss: {loss.item()}")
        loss += self.lmd * (outputs**2).mean()  # SD loss penalty
        logging.info(f"Training Step - SD loss: {loss.item()}")
        return loss

    def validation_step(self, batch):
        images, labels, _ = batch
        outputs = self.forward(images).squeeze()

        if outputs.shape == torch.Size([]):
            return

        print(f"Shape of outputs (validation): {outputs.shape}")
        print(f"Shape of labels (validation): {labels.shape}")

        loss = F.binary_cross_entropy_with_logits(outputs, labels.float())
        preds = torch.sigmoid(outputs)
        self.log("val_loss", loss, prog_bar=True, sync_dist=True)
        self.log(
            "val_acc",
            self.accuracy(preds, labels.int()),
            prog_bar=True,
            sync_dist=True,
        )
        self.log(
            "val_recall",
            self.recall(preds, labels.int()),
            prog_bar=True,
            sync_dist=True,
        )
        output = {"val_loss": loss, "preds": preds, "labels": labels}
        self.validation_outputs.append(output)
        logging.info(f"Validation Step - Batch loss: {loss.item()}")
        return output

    def predict_step(self, batch):
        images, label, domain = batch
        outputs = self.forward(images).squeeze()
        preds = torch.sigmoid(outputs)
        return preds, label, domain

    def on_validation_epoch_end(self):
        if not self.validation_outputs:
            logging.warning("No outputs in validation step to process")
            return
        preds = torch.cat([x["preds"] for x in self.validation_outputs])
        labels = torch.cat([x["labels"] for x in self.validation_outputs])
        if labels.unique().size(0) == 1:
            logging.warning("Only one class in validation step")
            return
        auc_score = roc_auc_score(labels.cpu(), preds.cpu())
        self.log("val_auc", auc_score, prog_bar=True, sync_dist=True)
        logging.info(f"Validation Epoch End - AUC score: {auc_score}")
        self.validation_outputs = []

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.0005)
        return optimizer


checkpoint_callback = ModelCheckpoint(
    monitor="val_loss",
    dirpath="./model_checkpoints/",
    filename="image-classifier-{step}-{val_loss:.2f}",
    save_top_k=3,
    mode="min",
    every_n_train_steps=1001,
    enable_version_counter=True,
)

early_stop_callback = EarlyStopping(
    monitor="val_loss",
    patience=4,
    mode="min",
)


def load_image(image_path, transform=None):
    image = Image.open(image_path).convert("RGB")

    if transform:
        image = transform(image)

    return image


def predict_single_image(image_path, model, transform=None):
    image = load_image(image_path, transform)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.to(device)

    image = image.to(device)

    model.eval()

    with torch.no_grad():
        image = image.unsqueeze(0)
        output = model(image).squeeze()
        print(output)
        prediction = torch.sigmoid(output).item()

    return prediction


parser = argparse.ArgumentParser()
parser.add_argument(
    "--ckpt_path", help="checkpoint to continue from", required=False
)
parser.add_argument(
    "--predict", help="predict on test set", action="store_true"
)
parser.add_argument("--reset", help="reset training", action="store_true")
parser.add_argument(
    "--predict_image",
    help="predict the class of a single image",
    action="store_true",
)
parser.add_argument(
    "--image_path",
    help="path to the image to predict",
    type=str,
    required=False,
)
args = parser.parse_args()

train_domains = [0, 1, 4]
val_domains = [0, 1, 4]
lmd_value = 0

if args.predict:
    test_dl = load_dataloader(
        [0, 1, 2, 3, 4], "test", batch_size=128, num_workers=1
    )
    model = ImageClassifier.load_from_checkpoint(args.ckpt_path)
    trainer = pl.Trainer()
    predictions = trainer.predict(model, dataloaders=test_dl)
    preds, labels, domains = zip(*predictions)
    preds = torch.cat(preds).cpu().numpy()
    labels = torch.cat(labels).cpu().numpy()
    domains = torch.cat(domains).cpu().numpy()
    print(preds.shape, labels.shape, domains.shape)
    df = pd.DataFrame({"preds": preds, "labels": labels, "domains": domains})
    filename = "preds-" + args.ckpt_path.split("/")[-1]
    df.to_csv(f"outputs/{filename}.csv", index=False)
elif args.predict_image:
    image_path = args.image_path
    model = ImageClassifier.load_from_checkpoint(args.ckpt_path)

    # Define the transformations for the image
    transform = transforms.Compose(
        [
            transforms.Resize((224, 224)),  # Image size expected by ResNet50
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            ),
        ]
    )

    prediction = predict_single_image(image_path, model, transform)
    print("prediction", prediction)

    # Output the prediction
    print(
        f"Prediction for {image_path}: {'Human' if prediction <= 0.001 else 'Generated'}"
    )
else:
    train_dl = load_dataloader(
        train_domains, "train", batch_size=128, num_workers=4
    )
    logging.info("Training dataloader loaded")
    val_dl = load_dataloader(val_domains, "val", batch_size=128, num_workers=4)
    logging.info("Validation dataloader loaded")

    if args.reset:
        model = ImageClassifier.load_from_checkpoint(args.ckpt_path)
    else:
        model = ImageClassifier(lmd=lmd_value)
    trainer = pl.Trainer(
        callbacks=[checkpoint_callback, early_stop_callback],
        max_steps=20000,
        val_check_interval=1000,
        check_val_every_n_epoch=None,
    )
    trainer.fit(
        model=model,
        train_dataloaders=train_dl,
        val_dataloaders=val_dl,
        ckpt_path=args.ckpt_path if not args.reset else None,
    )