| import os | |
| from typing import Union, Tuple | |
| import torchvision.transforms.functional as F | |
| from torch import Tensor | |
| from torchvision.transforms import transforms | |
| from auxiliary.settings import USE_CONFIDENCE_WEIGHTED_POOLING | |
| from auxiliary.utils import correct, rescale, scale | |
| from classes.core.Model import Model | |
| from classes.fc4.FC4 import FC4 | |
| class ModelFC4(Model): | |
| def __init__(self): | |
| super().__init__() | |
| self._network = FC4().to(self._device) | |
| def predict(self, img: Tensor, return_steps: bool = False) -> Union[Tensor, Tuple]: | |
| """ | |
| Performs inference on the input image using the FC4 method. | |
| @param img: the image for which an illuminant colour has to be estimated | |
| @param return_steps: whether or not to also return the per-patch estimates and confidence weights. When this | |
| flag is set to True, confidence-weighted pooling must be active) | |
| @return: the colour estimate as a Tensor. If "return_steps" is set to true, the per-path colour estimates and | |
| the confidence weights are also returned (used for visualizations) | |
| """ | |
| if USE_CONFIDENCE_WEIGHTED_POOLING: | |
| pred, rgb, confidence = self._network(img) | |
| if return_steps: | |
| return pred, rgb, confidence | |
| return pred | |
| return self._network(img) | |
| def optimize(self, img: Tensor, label: Tensor) -> float: | |
| self._optimizer.zero_grad() | |
| pred = self.predict(img) | |
| loss = self.get_loss(pred, label) | |
| loss.backward() | |
| self._optimizer.step() | |
| return loss.item() | |