third_party/DarkFeat/trainer.py

import os
import cv2
import time
import yaml
import torch
import datetime
from tensorboardX import SummaryWriter
import torchvision.transforms as tvf
import torch.nn as nn
import torch.nn.functional as F

from nets.geom import getK, getWarp, _grid_positions, getWarpNoValidate
from nets.loss import make_detector_loss, make_noise_score_map_loss
from nets.score import extract_kpts
from nets.multi_sampler import MultiSampler
from nets.noise_reliability_loss import MultiPixelAPLoss
from datasets.noise_simulator import NoiseSimulator
from nets.l2net import Quad_L2Net


class Trainer:
    def __init__(self, config, device, loader, job_name, start_cnt):
        self.config = config
        self.device = device
        self.loader = loader

        # tensorboard writer construction
        os.makedirs("./runs/", exist_ok=True)
        if job_name != "":
            self.log_dir = f"runs/{job_name}"
        else:
            self.log_dir = f'runs/{datetime.datetime.now().strftime("%m-%d-%H%M%S")}'

        self.writer = SummaryWriter(self.log_dir)
        with open(f"{self.log_dir}/config.yaml", "w") as f:
            yaml.dump(config, f)

        if config["network"]["input_type"] == "gray":
            self.model = eval(f'{config["network"]["model"]}(inchan=1)').to(device)
        elif (
            config["network"]["input_type"] == "rgb"
            or config["network"]["input_type"] == "raw-demosaic"
        ):
            self.model = eval(f'{config["network"]["model"]}(inchan=3)').to(device)
        elif config["network"]["input_type"] == "raw":
            self.model = eval(f'{config["network"]["model"]}(inchan=4)').to(device)
        else:
            raise NotImplementedError()

        # noise maker
        self.noise_maker = NoiseSimulator(device)

        # reliability map conv
        self.model.clf = nn.Conv2d(128, 2, kernel_size=1).cuda()

        # load model
        self.cnt = 0
        if start_cnt != 0:
            self.model.load_state_dict(
                torch.load(
                    f"{self.log_dir}/model_{start_cnt:06d}.pth", map_location=device
                )
            )
            self.cnt = start_cnt + 1

        # sampler
        sampler = MultiSampler(
            ngh=7,
            subq=-8,
            subd=1,
            pos_d=3,
            neg_d=5,
            border=16,
            subd_neg=-8,
            maxpool_pos=True,
        ).to(device)
        self.reliability_relitive_loss = MultiPixelAPLoss(sampler, nq=20).to(device)

        # optimizer and scheduler
        if self.config["training"]["optimizer"] == "SGD":
            self.optimizer = torch.optim.SGD(
                [
                    {
                        "params": self.model.parameters(),
                        "initial_lr": self.config["training"]["lr"],
                    }
                ],
                lr=self.config["training"]["lr"],
                momentum=self.config["training"]["momentum"],
                weight_decay=self.config["training"]["weight_decay"],
            )
        elif self.config["training"]["optimizer"] == "Adam":
            self.optimizer = torch.optim.Adam(
                [
                    {
                        "params": self.model.parameters(),
                        "initial_lr": self.config["training"]["lr"],
                    }
                ],
                lr=self.config["training"]["lr"],
                weight_decay=self.config["training"]["weight_decay"],
            )
        else:
            raise NotImplementedError()

        self.lr_scheduler = torch.optim.lr_scheduler.StepLR(
            self.optimizer,
            step_size=self.config["training"]["lr_step"],
            gamma=self.config["training"]["lr_gamma"],
            last_epoch=start_cnt,
        )
        for param_tensor in self.model.state_dict():
            print(param_tensor, "\t", self.model.state_dict()[param_tensor].size())

    def save(self, iter_num):
        torch.save(self.model.state_dict(), f"{self.log_dir}/model_{iter_num:06d}.pth")

    def load(self, path):
        self.model.load_state_dict(torch.load(path))

    def train(self):
        self.model.train()

        for epoch in range(2):
            for batch_idx, inputs in enumerate(self.loader):
                self.optimizer.zero_grad()
                t = time.time()

                # preprocess and add noise
                img0_ori, noise_img0_ori = self.preprocess_noise_pair(
                    inputs["img0"], self.cnt
                )
                img1_ori, noise_img1_ori = self.preprocess_noise_pair(
                    inputs["img1"], self.cnt
                )

                img0 = img0_ori.permute(0, 3, 1, 2).float().to(self.device)
                img1 = img1_ori.permute(0, 3, 1, 2).float().to(self.device)
                noise_img0 = noise_img0_ori.permute(0, 3, 1, 2).float().to(self.device)
                noise_img1 = noise_img1_ori.permute(0, 3, 1, 2).float().to(self.device)

                if self.config["network"]["input_type"] == "rgb":
                    # 3-channel rgb
                    RGB_mean = [0.485, 0.456, 0.406]
                    RGB_std = [0.229, 0.224, 0.225]
                    norm_RGB = tvf.Normalize(mean=RGB_mean, std=RGB_std)
                    img0 = norm_RGB(img0)
                    img1 = norm_RGB(img1)
                    noise_img0 = norm_RGB(noise_img0)
                    noise_img1 = norm_RGB(noise_img1)

                elif self.config["network"]["input_type"] == "gray":
                    # 1-channel
                    img0 = torch.mean(img0, dim=1, keepdim=True)
                    img1 = torch.mean(img1, dim=1, keepdim=True)
                    noise_img0 = torch.mean(noise_img0, dim=1, keepdim=True)
                    noise_img1 = torch.mean(noise_img1, dim=1, keepdim=True)
                    norm_gray0 = tvf.Normalize(mean=img0.mean(), std=img0.std())
                    norm_gray1 = tvf.Normalize(mean=img1.mean(), std=img1.std())
                    img0 = norm_gray0(img0)
                    img1 = norm_gray1(img1)
                    noise_img0 = norm_gray0(noise_img0)
                    noise_img1 = norm_gray1(noise_img1)

                elif self.config["network"]["input_type"] == "raw":
                    # 4-channel
                    pass

                elif self.config["network"]["input_type"] == "raw-demosaic":
                    # 3-channel
                    pass

                else:
                    raise NotImplementedError()

                desc0, score_map0, _, _ = self.model(img0)
                desc1, score_map1, _, _ = self.model(img1)

                conf0 = F.softmax(self.model.clf(torch.abs(desc0) ** 2.0), dim=1)[
                    :, 1:2
                ]
                conf1 = F.softmax(self.model.clf(torch.abs(desc1) ** 2.0), dim=1)[
                    :, 1:2
                ]

                noise_desc0, noise_score_map0, noise_at0, noise_att0 = self.model(
                    noise_img0
                )
                noise_desc1, noise_score_map1, noise_at1, noise_att1 = self.model(
                    noise_img1
                )

                noise_conf0 = F.softmax(
                    self.model.clf(torch.abs(noise_desc0) ** 2.0), dim=1
                )[:, 1:2]
                noise_conf1 = F.softmax(
                    self.model.clf(torch.abs(noise_desc1) ** 2.0), dim=1
                )[:, 1:2]

                cur_feat_size0 = torch.tensor(score_map0.shape[2:])
                cur_feat_size1 = torch.tensor(score_map1.shape[2:])

                desc0 = desc0.permute(0, 2, 3, 1)
                desc1 = desc1.permute(0, 2, 3, 1)
                score_map0 = score_map0.permute(0, 2, 3, 1)
                score_map1 = score_map1.permute(0, 2, 3, 1)
                noise_desc0 = noise_desc0.permute(0, 2, 3, 1)
                noise_desc1 = noise_desc1.permute(0, 2, 3, 1)
                noise_score_map0 = noise_score_map0.permute(0, 2, 3, 1)
                noise_score_map1 = noise_score_map1.permute(0, 2, 3, 1)
                conf0 = conf0.permute(0, 2, 3, 1)
                conf1 = conf1.permute(0, 2, 3, 1)
                noise_conf0 = noise_conf0.permute(0, 2, 3, 1)
                noise_conf1 = noise_conf1.permute(0, 2, 3, 1)

                r_K0 = getK(inputs["ori_img_size0"], cur_feat_size0, inputs["K0"]).to(
                    self.device
                )
                r_K1 = getK(inputs["ori_img_size1"], cur_feat_size1, inputs["K1"]).to(
                    self.device
                )

                pos0 = _grid_positions(
                    cur_feat_size0[0], cur_feat_size0[1], img0.shape[0]
                ).to(self.device)

                pos0_for_rel, pos1_for_rel, _ = getWarpNoValidate(
                    pos0,
                    inputs["rel_pose"].to(self.device),
                    inputs["depth0"].to(self.device),
                    r_K0,
                    inputs["depth1"].to(self.device),
                    r_K1,
                    img0.shape[0],
                )

                pos0, pos1, _ = getWarp(
                    pos0,
                    inputs["rel_pose"].to(self.device),
                    inputs["depth0"].to(self.device),
                    r_K0,
                    inputs["depth1"].to(self.device),
                    r_K1,
                    img0.shape[0],
                )

                reliab_loss_relative = self.reliability_relitive_loss(
                    desc0,
                    desc1,
                    noise_desc0,
                    noise_desc1,
                    conf0,
                    conf1,
                    noise_conf0,
                    noise_conf1,
                    pos0_for_rel,
                    pos1_for_rel,
                    img0.shape[0],
                    img0.shape[2],
                    img0.shape[3],
                )

                det_structured_loss, det_accuracy = make_detector_loss(
                    pos0,
                    pos1,
                    desc0,
                    desc1,
                    score_map0,
                    score_map1,
                    img0.shape[0],
                    self.config["network"]["use_corr_n"],
                    self.config["network"]["loss_type"],
                    self.config,
                )

                det_structured_loss_noise, det_accuracy_noise = make_detector_loss(
                    pos0,
                    pos1,
                    noise_desc0,
                    noise_desc1,
                    noise_score_map0,
                    noise_score_map1,
                    img0.shape[0],
                    self.config["network"]["use_corr_n"],
                    self.config["network"]["loss_type"],
                    self.config,
                )

                indices0, scores0 = extract_kpts(
                    score_map0.permute(0, 3, 1, 2),
                    k=self.config["network"]["det"]["kpt_n"],
                    score_thld=self.config["network"]["det"]["score_thld"],
                    nms_size=self.config["network"]["det"]["nms_size"],
                    eof_size=self.config["network"]["det"]["eof_size"],
                    edge_thld=self.config["network"]["det"]["edge_thld"],
                )
                indices1, scores1 = extract_kpts(
                    score_map1.permute(0, 3, 1, 2),
                    k=self.config["network"]["det"]["kpt_n"],
                    score_thld=self.config["network"]["det"]["score_thld"],
                    nms_size=self.config["network"]["det"]["nms_size"],
                    eof_size=self.config["network"]["det"]["eof_size"],
                    edge_thld=self.config["network"]["det"]["edge_thld"],
                )

                noise_score_loss0, mask0 = make_noise_score_map_loss(
                    score_map0, noise_score_map0, indices0, img0.shape[0], thld=0.1
                )
                noise_score_loss1, mask1 = make_noise_score_map_loss(
                    score_map1, noise_score_map1, indices1, img1.shape[0], thld=0.1
                )

                total_loss = det_structured_loss + det_structured_loss_noise
                total_loss += noise_score_loss0 / 2.0 * 1.0
                total_loss += noise_score_loss1 / 2.0 * 1.0
                total_loss += reliab_loss_relative[0] / 2.0 * 0.5
                total_loss += reliab_loss_relative[1] / 2.0 * 0.5

                self.writer.add_scalar("acc/normal_acc", det_accuracy, self.cnt)
                self.writer.add_scalar("acc/noise_acc", det_accuracy_noise, self.cnt)
                self.writer.add_scalar("loss/total_loss", total_loss, self.cnt)
                self.writer.add_scalar(
                    "loss/noise_score_loss",
                    (noise_score_loss0 + noise_score_loss1) / 2.0,
                    self.cnt,
                )
                self.writer.add_scalar(
                    "loss/det_loss_normal", det_structured_loss, self.cnt
                )
                self.writer.add_scalar(
                    "loss/det_loss_noise", det_structured_loss_noise, self.cnt
                )
                print(
                    "iter={},\tloss={:.4f},\tacc={:.4f},\t{:.4f}s/iter".format(
                        self.cnt, total_loss, det_accuracy, time.time() - t
                    )
                )
                # print(f'normal_loss: {det_structured_loss}, noise_loss: {det_structured_loss_noise}, reliab_loss: {reliab_loss_relative[0]}, {reliab_loss_relative[1]}')

                if det_structured_loss != 0:
                    total_loss.backward()
                    self.optimizer.step()
                self.lr_scheduler.step()

                if self.cnt % 100 == 0:
                    noise_indices0, noise_scores0 = extract_kpts(
                        noise_score_map0.permute(0, 3, 1, 2),
                        k=self.config["network"]["det"]["kpt_n"],
                        score_thld=self.config["network"]["det"]["score_thld"],
                        nms_size=self.config["network"]["det"]["nms_size"],
                        eof_size=self.config["network"]["det"]["eof_size"],
                        edge_thld=self.config["network"]["det"]["edge_thld"],
                    )
                    noise_indices1, noise_scores1 = extract_kpts(
                        noise_score_map1.permute(0, 3, 1, 2),
                        k=self.config["network"]["det"]["kpt_n"],
                        score_thld=self.config["network"]["det"]["score_thld"],
                        nms_size=self.config["network"]["det"]["nms_size"],
                        eof_size=self.config["network"]["det"]["eof_size"],
                        edge_thld=self.config["network"]["det"]["edge_thld"],
                    )
                    if self.config["network"]["input_type"] == "raw":
                        kpt_img0 = self.showKeyPoints(
                            img0_ori[0][..., :3] * 255.0, indices0[0]
                        )
                        kpt_img1 = self.showKeyPoints(
                            img1_ori[0][..., :3] * 255.0, indices1[0]
                        )
                        noise_kpt_img0 = self.showKeyPoints(
                            noise_img0_ori[0][..., :3] * 255.0, noise_indices0[0]
                        )
                        noise_kpt_img1 = self.showKeyPoints(
                            noise_img1_ori[0][..., :3] * 255.0, noise_indices1[0]
                        )
                    else:
                        kpt_img0 = self.showKeyPoints(img0_ori[0] * 255.0, indices0[0])
                        kpt_img1 = self.showKeyPoints(img1_ori[0] * 255.0, indices1[0])
                        noise_kpt_img0 = self.showKeyPoints(
                            noise_img0_ori[0] * 255.0, noise_indices0[0]
                        )
                        noise_kpt_img1 = self.showKeyPoints(
                            noise_img1_ori[0] * 255.0, noise_indices1[0]
                        )

                    self.writer.add_image(
                        "img0/kpts", kpt_img0, self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img1/kpts", kpt_img1, self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img0/noise_kpts", noise_kpt_img0, self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img1/noise_kpts", noise_kpt_img1, self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img0/score_map", score_map0[0], self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img1/score_map", score_map1[0], self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img0/noise_score_map",
                        noise_score_map0[0],
                        self.cnt,
                        dataformats="HWC",
                    )
                    self.writer.add_image(
                        "img1/noise_score_map",
                        noise_score_map1[0],
                        self.cnt,
                        dataformats="HWC",
                    )
                    self.writer.add_image(
                        "img0/kpt_mask", mask0.unsqueeze(2), self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img1/kpt_mask", mask1.unsqueeze(2), self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img0/conf", conf0[0], self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img1/conf", conf1[0], self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img0/noise_conf", noise_conf0[0], self.cnt, dataformats="HWC"
                    )
                    self.writer.add_image(
                        "img1/noise_conf", noise_conf1[0], self.cnt, dataformats="HWC"
                    )

                if self.cnt % 5000 == 0:
                    self.save(self.cnt)

                self.cnt += 1

    def showKeyPoints(self, img, indices):
        key_points = cv2.KeyPoint_convert(indices.cpu().float().numpy()[:, ::-1])
        img = img.numpy().astype("uint8")
        img = cv2.drawKeypoints(img, key_points, None, color=(0, 255, 0))
        return img

    def preprocess(self, img, iter_idx):
        if (
            not self.config["network"]["noise"]
            and "raw" not in self.config["network"]["input_type"]
        ):
            return img

        raw = self.noise_maker.rgb2raw(img, batched=True)

        if self.config["network"]["noise"]:
            ratio_dec = (
                min(self.config["network"]["noise_maxstep"], iter_idx)
                / self.config["network"]["noise_maxstep"]
            )
            raw = self.noise_maker.raw2noisyRaw(raw, ratio_dec=ratio_dec, batched=True)

        if self.config["network"]["input_type"] == "raw":
            return torch.tensor(self.noise_maker.raw2packedRaw(raw, batched=True))

        if self.config["network"]["input_type"] == "raw-demosaic":
            return torch.tensor(self.noise_maker.raw2demosaicRaw(raw, batched=True))

        rgb = self.noise_maker.raw2rgb(raw, batched=True)
        if (
            self.config["network"]["input_type"] == "rgb"
            or self.config["network"]["input_type"] == "gray"
        ):
            return torch.tensor(rgb)

        raise NotImplementedError()

    def preprocess_noise_pair(self, img, iter_idx):
        assert self.config["network"]["noise"]

        raw = self.noise_maker.rgb2raw(img, batched=True)

        ratio_dec = (
            min(self.config["network"]["noise_maxstep"], iter_idx)
            / self.config["network"]["noise_maxstep"]
        )
        noise_raw = self.noise_maker.raw2noisyRaw(
            raw, ratio_dec=ratio_dec, batched=True
        )

        if self.config["network"]["input_type"] == "raw":
            return torch.tensor(
                self.noise_maker.raw2packedRaw(raw, batched=True)
            ), torch.tensor(self.noise_maker.raw2packedRaw(noise_raw, batched=True))

        if self.config["network"]["input_type"] == "raw-demosaic":
            return torch.tensor(
                self.noise_maker.raw2demosaicRaw(raw, batched=True)
            ), torch.tensor(self.noise_maker.raw2demosaicRaw(noise_raw, batched=True))

        noise_rgb = self.noise_maker.raw2rgb(noise_raw, batched=True)
        if (
            self.config["network"]["input_type"] == "rgb"
            or self.config["network"]["input_type"] == "gray"
        ):
            return img, torch.tensor(noise_rgb)

        raise NotImplementedError()