import torch

import torch.nn as nn

from weakly_supervised_localization import WSLoss

class Interpreter:

    def __init__(self, cnn, cam, loader):

        self.cnn = cnn

        self.cam = cam

        self.loader = loader

        self.loss_fn = nn.BCEWithLogitsLoss()

        self.ws_loss_fn = WSLoss(self.loss_fn)

        self.optimizer = torch.optim.Adam(self.cnn.parameters())

    def train_weakly_supervised(self):

        for images, labels in self.loader:

            self.optimizer.zero_grad()

            preds = self.cnn(images)

            cam_maps = self.cam(images)

            loss = self.ws_loss_fn(preds, cam_maps, labels)

            loss.backward()

            self.optimizer.step()