# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/metrics.ipynb.

# %% auto 0
__all__ = ['SegmentationMetric', 'IOU_Score', 'Accuracy', 'Recall']

# %% ../nbs/metrics.ipynb 2
from torcheval.metrics.metric import Metric
import segmentation_models_pytorch as smp
import torch

# %% ../nbs/metrics.ipynb 3
class SegmentationMetric(Metric):
    def __init__(self, num_classes, device=None ) -> None:
        super().__init__(device=device)
        self._add_state('weighted_sum', torch.tensor(0.0, device=self.device, dtype=torch.float64))
        self._add_state('weights', torch.tensor(0.0, device=self.device, dtype=torch.float64))
        self.num_classes = num_classes
        
    @torch.inference_mode()
    def compute(self):
        return self.weighted_sum/self.weights
    
    @torch.inference_mode()
    def merge_state(self, metrics):
        for metric in metrics:
            self.weighted_sum += metric.weighted_sum.to(self.device)
            self.weights += metric.weights.to(self.device)
        return self


# %% ../nbs/metrics.ipynb 4
class IOU_Score(SegmentationMetric):
    def __init__(self, num_classes, device=None, reduction='micro') -> None:
        super().__init__(num_classes=num_classes, device=device)
        self.reduction = reduction
    
    @torch.inference_mode()
    def update(self, preds, target):
        if len(preds.shape)==4:
            preds = preds.softmax(dim=1).argmax(dim=1)
        tp, fp, fn, tn = smp.metrics.get_stats(preds, target, mode='multiclass', num_classes=self.num_classes)
        self.weighted_sum += smp.metrics.iou_score(tp, fp, fn, tn, reduction=self.reduction)
        self.weights += 1# tp.shape[0]
        return self
    

    
class Accuracy(SegmentationMetric):
    def __init__(self, num_classes, device=None, reduction='macro') -> None:
        super().__init__(num_classes=num_classes, device=device)
        self.reduction = reduction

    @torch.inference_mode()
    def update(self, preds, target):
        if len(preds.shape)==4:
            preds = preds.softmax(dim=1).argmax(dim=1)
        tp, fp, fn, tn = smp.metrics.get_stats(preds, target, mode='multiclass', num_classes=self.num_classes)
        self.weighted_sum += smp.metrics.accuracy(tp, fp, fn, tn, reduction=self.reduction)
        self.weights += 1# tp.shape[0]
        return self


class Recall(SegmentationMetric):
    def __init__(self, num_classes, device=None, reduction='micro-imagewise') -> None:
        super().__init__(num_classes=num_classes, device=device)
        self.reduction = reduction

    @torch.inference_mode()
    def update(self, preds, target):
        if len(preds.shape)==4:
            preds = preds.softmax(dim=1).argmax(dim=1)
        tp, fp, fn, tn = smp.metrics.get_stats(preds, target, mode='multiclass', num_classes=self.num_classes)
        self.weighted_sum += smp.metrics.recall(tp, fp, fn, tn, reduction=self.reduction)
        self.weights += 1# tp.shape[0]
        return self

