MiAlgo

MiAlgo/Dockerfile

@@ -0,0 +1,9 @@
+FROM q935970314/mialgo:2022
+
+COPY . /workdir/
+
+RUN export PATH=/opt/conda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin && pip install grayness_index_python
+
+WORKDIR /workdir/
+
+

MiAlgo/README.txt

@@ -0,0 +1,16 @@
+Run the code to build docker and process the night image:
+
+docker build -t mialgo24 .
+
+cd DATA_FOLDER_PATH
+docker run -it --rm --gpus all -v $(pwd)/data:/data mialgo24 ./run.sh
+
+
+Please note that the "--gpus all" parameter is used and ensure that the model is running on the GPU.
+We first load all images into memory. If the memory explodes, please process less than 50 images at a time.
+
+We will open source the code after the competition ends.
+
+If you have any questions, please contact me in time, thank you.
+
+Email:[email protected]

MiAlgo/assets/pretrained/hub/checkpoints/squeezenet1_1-f364aa15.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f364aa15cc776cd43e679ca0859f479db11ef4852a4e79bb237519f9d16617c5
+size 4966400

MiAlgo/auxiliary/settings.py

@@ -0,0 +1,47 @@
+import re
+
+import numpy as np
+import torch
+
+
+# --- Determinism (for reproducibility) ---
+
+def make_deterministic(seed: int):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    torch.backends.cudnn.benchmark = False
+
+
+# --- Device (cpu or cuda:n) ---
+
+DEVICE_TYPE = "cuda:0"
+
+
+def get_device() -> torch.device:
+    if DEVICE_TYPE == "cpu":
+        print("\n Running on device 'cpu' \n")
+        return torch.device("cpu")
+
+    if re.match(r"\bcuda:\b\d+", DEVICE_TYPE):
+        if not torch.cuda.is_available():
+            print("\n WARNING: running on cpu since device {} is not available \n".format(DEVICE_TYPE))
+            return torch.device("cpu")
+
+        # print("\n Running on device '{}' \n".format(DEVICE_TYPE))
+        return torch.device(DEVICE_TYPE)
+
+    raise ValueError("ERROR: {} is not a valid device! Supported device are 'cpu' and 'cuda:n'".format(DEVICE_TYPE))
+
+
+DEVICE = get_device()
+
+# --- Model ---
+
+# If set to False, a simpler summation pooling will be used
+USE_CONFIDENCE_WEIGHTED_POOLING = True
+if not USE_CONFIDENCE_WEIGHTED_POOLING:
+    print("\n WARN: confidence-weighted pooling option is set to False \n")
+
+# Input size
+TRAIN_IMG_W, TRAIN_IMG_H = 512, 512
+TEST_IMG_W, TEST_IMG_H = 0, 0

MiAlgo/auxiliary/utils.py

@@ -0,0 +1,107 @@
+import math
+import os
+from typing import Union, List, Tuple
+
+import numpy as np
+import torch
+import torchvision.transforms.functional as F
+from PIL.Image import Image
+from scipy.spatial.distance import jensenshannon
+from torch import Tensor
+from torch.nn.functional import interpolate
+
+from auxiliary.settings import DEVICE
+
+
+
+def print_metrics(current_metrics: dict, best_metrics: dict):
+    print(" Mean ......... : {:.4f} (Best: {:.4f})".format(current_metrics["mean"], best_metrics["mean"]))
+    print(" Median ....... : {:.4f} (Best: {:.4f})".format(current_metrics["median"], best_metrics["median"]))
+    print(" Trimean ...... : {:.4f} (Best: {:.4f})".format(current_metrics["trimean"], best_metrics["trimean"]))
+    print(" Best 25% ..... : {:.4f} (Best: {:.4f})".format(current_metrics["bst25"], best_metrics["bst25"]))
+    print(" Worst 25% .... : {:.4f} (Best: {:.4f})".format(current_metrics["wst25"], best_metrics["wst25"]))
+    print(" Worst 5% ..... : {:.4f} (Best: {:.4f})".format(current_metrics["wst5"], best_metrics["wst5"]))
+
+
+def correct(img: Image, illuminant: Tensor) -> Image:
+    """
+    Corrects the color of the illuminant of a linear image based on an estimated (linear) illuminant
+    @param img: a linear image
+    @param illuminant: a linear illuminant
+    @return: a non-linear color-corrected version of the input image
+    """
+    img = F.to_tensor(img).to(DEVICE)
+
+    # Correct the image
+    correction = illuminant.unsqueeze(2).unsqueeze(3) * torch.sqrt(Tensor([3])).to(DEVICE)
+    corrected_img = torch.div(img, correction + 1e-10)
+
+    # Normalize the image
+    max_img = torch.max(torch.max(torch.max(corrected_img, dim=1)[0], dim=1)[0], dim=1)[0] + 1e-10
+    max_img = max_img.unsqueeze(1).unsqueeze(1).unsqueeze(1)
+    normalized_img = torch.div(corrected_img, max_img)
+
+    return F.to_pil_image(linear_to_nonlinear(normalized_img).squeeze(), mode="RGB")
+
+
+def linear_to_nonlinear(img: Union[np.array, Image, Tensor]) -> Union[np.array, Image, Tensor]:
+    if isinstance(img, np.ndarray):
+        return np.power(img, (1.0 / 2.2))
+    if isinstance(img, Tensor):
+        return torch.pow(img, 1.0 / 2.2)
+    return F.to_pil_image(torch.pow(F.to_tensor(img), 1.0 / 2.2).squeeze(), mode="RGB")
+
+
+def normalize(img: np.ndarray) -> np.ndarray:
+    max_int = 65535.0
+    return np.clip(img, 0.0, max_int) * (1.0 / max_int)
+
+
+def rgb_to_bgr(x: np.ndarray) -> np.ndarray:
+    return x[::-1]
+
+
+def bgr_to_rgb(x: np.ndarray) -> np.ndarray:
+    return x[:, :, ::-1]
+
+
+def hwc_to_chw(x: np.ndarray) -> np.ndarray:
+    """ Converts an image from height x width x channels to channels x height x width """
+    return x.transpose(2, 0, 1)
+
+
+def scale(x: Tensor) -> Tensor:
+    """ Scales all values of a tensor between 0 and 1 """
+    x = x - x.min()
+    x = x / x.max()
+    return x
+
+
+def rescale(x: Tensor, size: Tuple) -> Tensor:
+    """ Rescale tensor to image size for better visualization """
+    return interpolate(x, size, mode='bilinear')
+
+
+def angular_error(x: Tensor, y: Tensor, safe_v: float = 0.999999) -> Tensor:
+    x, y = torch.nn.functional.normalize(x, dim=1), torch.nn.functional.normalize(y, dim=1)
+    dot = torch.clamp(torch.sum(x * y, dim=1), -safe_v, safe_v)
+    angle = torch.acos(dot) * (180 / math.pi)
+    return torch.mean(angle).item()
+
+
+def tvd(pred: Tensor, label: Tensor) -> Tensor:
+    """
+    Total Variation Distance (TVD) is a distance measure for probability distributions
+    https://en.wikipedia.org/wiki/Total_variation_distance_of_probability_measures
+    """
+    return (Tensor([0.5]) * torch.abs(pred - label)).sum()
+
+
+def jsd(p: List, q: List) -> float:
+    """
+    Jensen-Shannon Divergence (JSD) between two probability distributions as square of scipy's JS distance. Refs:
+    - https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.jensenshannon.html
+    - https://stackoverflow.com/questions/15880133/jensen-shannon-divergence
+    """
+    return jensenshannon(p, q) ** 2
+

MiAlgo/checkpoint/nn_enhance.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37c37648922f149ebefcec8b4c32c687bbafacd1efd99e288b235aab6e834ad5
+size 152303133

MiAlgo/checkpoint/raw_denoise.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a987213026b43e214196154494e5b54045721dea11dcb8e6f1f2d8e632445699
+size 2072446

MiAlgo/classes/core/Evaluator.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+
+class Evaluator:
+
+    def __init__(self):
+        monitored_metrics = ["mean", "median", "trimean", "bst25", "wst25", "wst5"]
+        self.__metrics = {}
+        self.__best_metrics = {m: 100.0 for m in monitored_metrics}
+        self.__errors = []
+
+    def add_error(self, error: float):
+        self.__errors.append(error)
+
+    def reset_errors(self):
+        self.__errors = []
+
+    def get_errors(self) -> list:
+        return self.__errors
+
+    def get_metrics(self) -> dict:
+        return self.__metrics
+
+    def get_best_metrics(self) -> dict:
+        return self.__best_metrics
+
+    def compute_metrics(self) -> dict:
+        self.__errors = sorted(self.__errors)
+        self.__metrics = {
+            "mean": np.mean(self.__errors),
+            "median": self.__g(0.5),
+            "trimean": 0.25 * (self.__g(0.25) + 2 * self.__g(0.5) + self.__g(0.75)),
+            "bst25": np.mean(self.__errors[:int(0.25 * len(self.__errors))]),
+            "wst25": np.mean(self.__errors[int(0.75 * len(self.__errors)):]),
+            "wst5": self.__g(0.95)
+        }
+        return self.__metrics
+
+    def update_best_metrics(self) -> dict:
+        self.__best_metrics["mean"] = self.__metrics["mean"]
+        self.__best_metrics["median"] = self.__metrics["median"]
+        self.__best_metrics["trimean"] = self.__metrics["trimean"]
+        self.__best_metrics["bst25"] = self.__metrics["bst25"]
+        self.__best_metrics["wst25"] = self.__metrics["wst25"]
+        self.__best_metrics["wst5"] = self.__metrics["wst5"]
+        return self.__best_metrics
+
+    def __g(self, f: float) -> float:
+        return np.percentile(self.__errors, f * 100)

MiAlgo/classes/core/Loss.py

@@ -0,0 +1,13 @@
+import torch
+from torch import Tensor
+
+
+class Loss:
+    def __init__(self, device: torch.device):
+        self._device = device
+
+    def _compute(self, *args, **kwargs) -> Tensor:
+        pass
+
+    def __call__(self, *args, **kwargs):
+        return self._compute(*args).to(self._device)

MiAlgo/classes/core/LossTracker.py

@@ -0,0 +1,16 @@
+class LossTracker(object):
+
+    def __init__(self):
+        self.val, self.avg, self.sum, self.count = 0, 0, 0, 0
+
+    def reset(self):
+        self.val, self.avg, self.sum, self.count = 0, 0, 0, 0
+
+    def update(self, val: float, n: int = 1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def get_loss(self) -> float:
+        return self.avg

MiAlgo/classes/core/Model.py

@@ -0,0 +1,43 @@
+import os
+
+import torch
+from torch import Tensor
+
+from auxiliary.settings import DEVICE
+from classes.losses.AngularLoss import AngularLoss
+
+
+class Model:
+    def __init__(self):
+        self._device = DEVICE
+        self._criterion = AngularLoss(self._device)
+        self._optimizer = None
+        self._network = None
+
+    def print_network(self):
+        print("\n----------------------------------------------------------\n")
+        print(self._network)
+        print("\n----------------------------------------------------------\n")
+
+    def log_network(self, path_to_log: str):
+        open(os.path.join(path_to_log, "network.txt"), 'a+').write(str(self._network))
+
+    def get_loss(self, pred: Tensor, label: Tensor) -> Tensor:
+        return self._criterion(pred, label)
+
+    def train_mode(self):
+        self._network = self._network.train()
+
+    def evaluation_mode(self):
+        self._network = self._network.eval()
+
+    def save(self, path_to_log: str):
+        torch.save(self._network.state_dict(), os.path.join(path_to_log, "model.pth"))
+
+    def load(self, path_to_pretrained: str):
+        path_to_model = os.path.join(path_to_pretrained, "model.pth")
+        self._network.load_state_dict(torch.load(path_to_model, map_location=self._device))
+
+    def set_optimizer(self, learning_rate: float, optimizer_type: str = "adam"):
+        optimizers_map = {"adam": torch.optim.Adam, "rmsprop": torch.optim.RMSprop}
+        self._optimizer = optimizers_map[optimizer_type](self._network.parameters(), lr=learning_rate)

MiAlgo/classes/data/ColorCheckerDataset.py

@@ -0,0 +1,52 @@
+import os
+from typing import Tuple
+
+import numpy as np
+import scipy.io
+import torch
+import torch.utils.data as data
+
+from auxiliary.utils import normalize, bgr_to_rgb, linear_to_nonlinear, hwc_to_chw
+from classes.data.DataAugmenter import DataAugmenter
+
+
+class ColorCheckerDataset(data.Dataset):
+
+    def __init__(self, train: bool = True, folds_num: int = 1):
+
+        self.__train = train
+        self.__da = DataAugmenter()
+
+        path_to_folds = os.path.join("dataset", "folds.mat")
+        path_to_metadata = os.path.join("dataset", "metadata.txt")
+        self.__path_to_data = os.path.join("dataset", "preprocessed", "numpy_data")
+        self.__path_to_label = os.path.join("dataset", "preprocessed", "numpy_labels")
+
+        folds = scipy.io.loadmat(path_to_folds)
+        img_idx = folds["tr_split" if self.__train else "te_split"][0][folds_num][0]
+
+        metadata = open(path_to_metadata, 'r').readlines()
+        self.__fold_data = [metadata[i - 1] for i in img_idx]
+
+    def __getitem__(self, index: int) -> Tuple:
+        file_name = self.__fold_data[index].strip().split(' ')[1]
+        img = np.array(np.load(os.path.join(self.__path_to_data, file_name + '.npy')), dtype='float32')
+        illuminant = np.array(np.load(os.path.join(self.__path_to_label, file_name + '.npy')), dtype='float32')
+
+        if self.__train:
+            img, illuminant = self.__da.augment(img, illuminant)
+        else:
+            img = self.__da.crop(img)
+
+        img = hwc_to_chw(linear_to_nonlinear(bgr_to_rgb(normalize(img))))
+
+        img = torch.from_numpy(img.copy())
+        illuminant = torch.from_numpy(illuminant.copy())
+
+        if not self.__train:
+            img = img.type(torch.FloatTensor)
+
+        return img, illuminant, file_name
+
+    def __len__(self) -> int:
+        return len(self.__fold_data)

MiAlgo/classes/data/DataAugmenter.py

@@ -0,0 +1,152 @@
+import math
+import random
+
+import cv2
+import numpy as np
+
+from auxiliary.settings import TRAIN_IMG_H, TRAIN_IMG_W, TEST_IMG_H, TEST_IMG_W
+from auxiliary.utils import rgb_to_bgr
+
+
+class DataAugmenter:
+
+    def __init__(self):
+        # Input Size of the fully-convolutional network (SqueezeNet)
+        self.__train_size = (TRAIN_IMG_H, TRAIN_IMG_W)
+        self.__test_size = (TEST_IMG_H, TEST_IMG_W)
+
+        # Rotation angle
+        self.__angle = 60
+
+        # Patch scale
+        self.__scale = [0.1, 1.0]
+
+        # Color rescaling
+        self.__color = 0.8
+
+    @staticmethod
+    def __rotate_image(image: np.ndarray, angle: float) -> np.ndarray:
+        """
+        Rotates an OpenCV 2 / NumPy image about it's centre by the given angle (in degrees).
+        The returned image will be large enough to hold the entire new image, with a black background
+        """
+
+        # Get the image size (note: NumPy stores image matrices backwards)
+        image_size = (image.shape[1], image.shape[0])
+        image_center = tuple(np.array(image_size) / 2)
+
+        # Convert the OpenCV 3x2 rotation matrix to 3x3
+        rot_mat = np.vstack([cv2.getRotationMatrix2D(image_center, angle, 1.0), [0, 0, 1]])
+
+        rot_mat_no_translate = np.matrix(rot_mat[0:2, 0:2])
+
+        image_w2, image_h2 = image_size[0] * 0.5, image_size[1] * 0.5
+
+        # Obtain the rotated coordinates of the image corners
+        rotated_coords = [
+            (np.array([-image_w2, image_h2]) * rot_mat_no_translate).A[0],
+            (np.array([image_w2, image_h2]) * rot_mat_no_translate).A[0],
+            (np.array([-image_w2, -image_h2]) * rot_mat_no_translate).A[0],
+            (np.array([image_w2, -image_h2]) * rot_mat_no_translate).A[0]
+        ]
+
+        # Find the size of the new image
+        x_coords = [pt[0] for pt in rotated_coords]
+        x_pos, x_neg = [x for x in x_coords if x > 0], [x for x in x_coords if x < 0]
+
+        y_coords = [pt[1] for pt in rotated_coords]
+        y_pos, y_neg = [y for y in y_coords if y > 0], [y for y in y_coords if y < 0]
+
+        right_bound, left_bound, top_bound, bot_bound = max(x_pos), min(x_neg), max(y_pos), min(y_neg)
+        new_w, new_h = int(abs(right_bound - left_bound)), int(abs(top_bound - bot_bound))
+
+        # We require a translation matrix to keep the image centred
+        trans_mat = np.matrix([[1, 0, int(new_w * 0.5 - image_w2)], [0, 1, int(new_h * 0.5 - image_h2)], [0, 0, 1]])
+
+        # Compute the transform for the combined rotation and translation
+        affine_mat = (np.matrix(trans_mat) * np.matrix(rot_mat))[0:2, :]
+
+        # Apply the transform
+        return cv2.warpAffine(image, affine_mat, (new_w, new_h), flags=cv2.INTER_LINEAR)
+
+    @staticmethod
+    def __largest_rotated_rect(w: float, h: float, angle: float) -> tuple:
+        """
+        Given a rectangle of size w x h that has been rotated by 'angle' (in radians), computes the width and height of
+        the largest possible axis-aligned rectangle within the rotated rectangle.
+
+        Original JS code by 'Andri' and Magnus Hoff from Stack Overflow. Converted to Python by Aaron Snoswell
+        """
+        quadrant = int(math.floor(angle / (math.pi / 2))) & 3
+        sign_alpha = angle if ((quadrant & 1) == 0) else math.pi - angle
+        alpha = (sign_alpha % math.pi + math.pi) % math.pi
+
+        bb_w = w * math.cos(alpha) + h * math.sin(alpha)
+        bb_h = w * math.sin(alpha) + h * math.cos(alpha)
+
+        length = h if (w < h) else w
+        d = length * math.cos(alpha)
+
+        gamma = math.atan2(bb_w, bb_w) if (w < h) else math.atan2(bb_w, bb_w)
+        delta = math.pi - alpha - gamma
+        a = d * math.sin(alpha) / math.sin(delta)
+
+        y = a * math.cos(gamma)
+        x = y * math.tan(gamma)
+
+        return bb_w - 2 * x, bb_h - 2 * y
+
+    def __crop_around_center(self, image: np.ndarray, width: float, height: float) -> np.ndarray:
+        """ Given a NumPy / OpenCV 2 image, crops it to the given width and height around it's centre point """
+
+        image_size = (image.shape[1], image.shape[0])
+        image_center = (int(image_size[0] * 0.5), int(image_size[1] * 0.5))
+
+        width = image_size[0] if width > image_size[0] else width
+        height = image_size[1] if height > image_size[1] else height
+
+        x1, x2 = int(image_center[0] - width * 0.5), int(image_center[0] + width * 0.5)
+        y1, y2 = int(image_center[1] - height * 0.5), int(image_center[1] + height * 0.5)
+
+        return cv2.resize(image[y1:y2, x1:x2], self.__train_size)
+
+    def __rotate_and_crop(self, image: np.ndarray, angle: float) -> np.ndarray:
+        width, height = image.shape[:2]
+        target_width, target_height = self.__largest_rotated_rect(width, height, math.radians(angle))
+        return self.__crop_around_center(self.__rotate_image(image, angle), target_width, target_height)
+
+    @staticmethod
+    def __random_flip(img: np.ndarray) -> np.ndarray:
+        """ Perform random left/right flip with probability 0.5 """
+        if random.randint(0, 1):
+            img = img[:, ::-1]
+        return img.astype(np.float32)
+
+    def augment(self, img: np.ndarray, illumination: np.ndarray) -> tuple:
+        scale = math.exp(random.random() * math.log(self.__scale[1] / self.__scale[0])) * self.__scale[0]
+        s = min(max(int(round(min(img.shape[:2]) * scale)), 10), min(img.shape[:2]))
+
+        start_x = random.randrange(0, img.shape[0] - s + 1)
+        start_y = random.randrange(0, img.shape[1] - s + 1)
+        img = img[start_x:start_x + s, start_y:start_y + s]
+
+        img = self.__rotate_and_crop(img, angle=(random.random() - 0.5) * self.__angle)
+        img = self.__random_flip(img)
+
+        color_aug = np.zeros(shape=(3, 3))
+        for i in range(3):
+            color_aug[i, i] = 1 + random.random() * self.__color - 0.5 * self.__color
+        img *= np.array([[[color_aug[0][0], color_aug[1][1], color_aug[2][2]]]], dtype=np.float32)
+        new_image = np.clip(img, 0, 65535)
+
+        new_illuminant = np.zeros_like(illumination)
+        illumination = rgb_to_bgr(illumination)
+        for i in range(3):
+            for j in range(3):
+                new_illuminant[i] += illumination[j] * color_aug[i, j]
+        new_illuminant = rgb_to_bgr(np.clip(new_illuminant, 0.01, 100))
+
+        return new_image, new_illuminant
+
+    def crop(self, img: np.ndarray, scale: float = 0.5) -> np.ndarray:
+        return cv2.resize(img, self.__test_size, fx=scale, fy=scale)

MiAlgo/classes/fc4/FC4.py

@@ -0,0 +1,63 @@
+from typing import Union
+
+import torch
+from torch import nn, Tensor
+from torch.nn.functional import normalize
+
+from auxiliary.settings import USE_CONFIDENCE_WEIGHTED_POOLING
+from classes.fc4.squeezenet.SqueezeNetLoader import SqueezeNetLoader
+
+"""
+FC4: Fully Convolutional Color Constancy with Confidence-weighted Pooling
+* Original code: https://github.com/yuanming-hu/fc4
+* Paper: https://www.microsoft.com/en-us/research/publication/fully-convolutional-color-constancy-confidence-weighted-pooling/
+"""
+
+
+class FC4(torch.nn.Module):
+
+    def __init__(self, squeezenet_version: float = 1.1):
+        super().__init__()
+
+        # SqueezeNet backbone (conv1-fire8) for extracting semantic features
+        squeezenet = SqueezeNetLoader(squeezenet_version).load(pretrained=True)
+        self.backbone = nn.Sequential(*list(squeezenet.children())[0][:12])
+
+        # Final convolutional layers (conv6 and conv7) to extract semi-dense feature maps
+        self.final_convs = nn.Sequential(
+            nn.MaxPool2d(kernel_size=2, stride=1, ceil_mode=True),
+            nn.Conv2d(512, 64, kernel_size=6, stride=1, padding=3),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.5),
+            nn.Conv2d(64, 4 if USE_CONFIDENCE_WEIGHTED_POOLING else 3, kernel_size=1, stride=1),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x: Tensor) -> Union[tuple, Tensor]:
+        """
+        Estimate an RGB colour for the illuminant of the input image
+        @param x: the image for which the colour of the illuminant has to be estimated
+        @return: the colour estimate as a Tensor. If confidence-weighted pooling is used, the per-path colour estimates
+        and the confidence weights are returned as well (used for visualizations)
+        """
+
+        x = self.backbone(x)
+        out = self.final_convs(x)
+
+        # Confidence-weighted pooling: "out" is a set of semi-dense feature maps
+        if USE_CONFIDENCE_WEIGHTED_POOLING:
+            # Per-patch color estimates (first 3 dimensions)
+            rgb = normalize(out[:, :3, :, :], dim=1)
+
+            # Confidence (last dimension)
+            confidence = out[:, 3:4, :, :]
+
+            # Confidence-weighted pooling
+            pred = normalize(torch.sum(torch.sum(rgb * confidence, 2), 2), dim=1)
+
+            return pred, rgb, confidence
+
+        # Summation pooling
+        pred = normalize(torch.sum(torch.sum(out, 2), 2), dim=1)
+
+        return pred

MiAlgo/classes/fc4/ModelFC4.py

@@ -0,0 +1,44 @@
+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()
+
+