What's a better way to compute the ROC on my test set?

[davidefiocco]

I am training a (binary) image classifier roughly as follows:

from flash import Trainer
from flash.image import ImageClassificationData, ImageClassifier
from pytorch_lightning.loggers import MLFlowLogger

datamodule = ImageClassificationData.from_folders(
    train_folder=train_folder,
    val_folder=val_folder,
    test_folder=test_folder,
    predict_folder=test_folder
)

mlf_logger = MLFlowLogger(
  experiment_name="my_experiment", tracking_uri="databricks"
)

trainer = Trainer(
    max_epochs=5
    logger=mlf_logger,
)

trainer.finetune(model, datamodule=datamodule, strategy="freeze")

Afterwards, I would like to

1. Compute predictions on the test set
2. Calculate the ROC on my test set, using ground-truth (binary) labels
3. Log with mlflow

For this, I came up with the very cumbersome

predictions = trainer.predict(model, datamodule=datamodule, return_predictions=True)

y_preds_list = []
y_targets_list = []
y_filenames_list = []
for i in range(len(predictions)):
    for j in range(len(predictions[i])):
        y_preds_list.append(
            torch.softmax(torch.unsqueeze(predictions[i][j]["preds"], dim=0), dim=1)
        )
        y_targets_list.append(predictions[i][j]["target"])
        y_filenames_list.append(predictions[i][j]["metadata"]["filepath"])

y_preds = torch.stack(y_preds_list, dim=1).squeeze()[:, 1]
y_targets = torch.stack(y_targets_list)

preds_df = pd.DataFrame(
    {
        "y_pred": y_preds.cpu().numpy(),
        "y_target": y_targets.cpu().numpy(),
        "y_filename": y_filenames_list,
    }
)

preds_df.to_csv("preds.csv", index=False)
mlf_logger.experiment.log_artifact(mlf_logger.run_id, "preds.csv")

from torchmetrics.classification import ROC

roc = ROC(task="binary", pos_label=1)
fpr, tpr, thresholds = roc(y_preds, y_targets)

roc_df = pd.DataFrame(
    {
        "fpr": fpr.cpu().numpy(),
        "tpr": tpr.cpu().numpy(),
        "thresholds": thresholds.cpu().numpy(),
    }
)
roc_df.to_csv("roc.csv", index=False)
mlf_logger.experiment.log_artifact(mlf_logger.run_id, "roc.csv")

Which works, but is very cumbersome and could likely be simplified...
The clumsiest parts above are those that deal with getting (binary) scores and ground-truth target labels.

Can someone suggest a way?