evaluation.py

import torch
import numpy as np
import cv2
from utils import parser
from models.TriGeoNet.TriGeoNet import TriGeoNet


if __name__ == "__main__":
    model_config = parser()

    left_image_path = r'xxxxx'
    right_image_path = r'xxxxx'
    model_config.ckpt = r'xxxxx'

    def read_image(path):
        img = cv2.imread(path, cv2.IMREAD_UNCHANGED).astype('float32')

        if len(img.shape) == 3:
            img = np.moveaxis(img, -1, 0) / 127.5 - 1.0
            model_config.maxdisp = 96
            model_config.mindisp = -96
            model_config.channels = 3
        else:
            img = np.expand_dims(img, axis=0)
            img = (img - np.mean(img)) / np.std(img)
            model_config.maxdisp = 64
            model_config.mindisp = -128
            model_config.channels = 1

        return img

    left = read_image(left_image_path)
    right = read_image(right_image_path)

    model = TriGeoNet(model_config)
    model.load_state_dict(torch.load(model_config.ckpt, map_location='cpu')['state_dict'])
    model = model.cuda()
    left = torch.from_numpy(left).unsqueeze(0).cuda()
    right = torch.from_numpy(right).unsqueeze(0).cuda()

    model.eval()
    with torch.no_grad():
        with torch.cuda.amp.autocast(enabled=True):
            disp4, disp3, disp2, disp1 = model(left, right)
            disp_result = np.squeeze(disp1.detach().cpu().numpy())
    
    cv2.imwrite('disp_result.tif', disp_result)