File size: 1,602 Bytes
82567db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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()
|