41 lines
1.6 KiB
Python
41 lines
1.6 KiB
Python
from enum import Enum
|
|
from operator import itemgetter
|
|
from typing import Mapping, Iterable
|
|
|
|
import numpy as np
|
|
from funcy import rcompose, rpartial
|
|
|
|
from image_prediction.exceptions import UnexpectedLabelFormat
|
|
from image_prediction.label_mapper.mapper import LabelMapper
|
|
|
|
|
|
class ProbabilityMapperKeys(Enum):
|
|
LABEL = "label"
|
|
PROBABILITIES = "probabilities"
|
|
|
|
|
|
|
|
class ProbabilityMapper(LabelMapper):
|
|
def __init__(self, labels: Mapping[int, str]):
|
|
self.__labels = labels
|
|
# String conversion in the middle due to floating point precision issues.
|
|
# See: https://stackoverflow.com/questions/56820/round-doesnt-seem-to-be-rounding-properly
|
|
self.__rounder = rcompose(rpartial(round, 4), str, float)
|
|
|
|
def __validate_array_label_format(self, probabilities: np.ndarray) -> None:
|
|
if not len(probabilities) == len(self.__labels):
|
|
raise UnexpectedLabelFormat(
|
|
f"Received fewer probabilities ({len(probabilities)}) than labels were passed ({len(self.__labels)})."
|
|
)
|
|
|
|
def __map_array(self, probabilities: np.ndarray) -> dict:
|
|
self.__validate_array_label_format(probabilities)
|
|
cls2prob = dict(
|
|
sorted(zip(self.__labels, list(map(self.__rounder, probabilities))), key=itemgetter(1), reverse=True)
|
|
)
|
|
most_likely = [*cls2prob][0]
|
|
return {ProbabilityMapperKeys.LABEL: most_likely, ProbabilityMapperKeys.PROBABILITIES: cls2prob}
|
|
|
|
def map_labels(self, probabilities: Iterable[np.ndarray]) -> Iterable[dict]:
|
|
return map(self.__map_array, probabilities)
|