Test.py

import os
import numpy as np
from scipy.ndimage import binary_erosion
import cv2
import matplotlib.pyplot as plt
import time

import tools as tb
from libs import ge, rectify
from common import Camera, PointCloud

def compute_disparity_map(img1Rect, img2Rect, min_disp, max_disp, save=False):
    # Compute disparity map
    start = time.time()
    num_disp = int(2 ** np.ceil(np.log2(max_disp - min_disp)))
    print(f'Number of disparities: {num_disp}')
    stereo = cv2.StereoSGBM.create(numDisparities=num_disp, minDisparity=min_disp,
                                   blockSize=3, uniquenessRatio=30,
                                   speckleWindowSize=100, speckleRange=5,
                                   disp12MaxDiff=1, P1=8 * 3 * 3 ** 2, P2=32 * 3 * 3 ** 2)

    # Add zero padding to the images
    pad1 = np.zeros((h, max(0, num_disp + min_disp)), dtype=img1Rect.dtype)
    pad2 = np.zeros((h, max(0, -min_disp)), dtype=img1Rect.dtype)
    img1Rect = np.hstack((pad1, img1Rect, pad2))
    img2Rect = np.hstack((pad1, img2Rect, pad2))
    disparity = stereo.compute(img1Rect, img2Rect)
    disparity = disparity.astype(float) / 16.0

    # Remove disparities computed from black pixels in the rectified images (those pixels add high values to the disparity map)
    disparity[img2Rect == 0] = min_disp - 1
    disparity[img1Rect == 0] = min_disp - 1

    # Remove padding
    disparity = disparity[:, pad1.shape[1]: -pad2.shape[1]]
    plot_disparity_map(disparity)

    # Apply erosion to remove noise
    # Define the structuring element (vertical kernel for horizontal line removal)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 15))  # 1-pixel wide, 15 pixels tall

    # Apply erosion
    disparity = cv2.erode(disparity, kernel, iterations=1)
    # Optional: Apply dilation to reconstruct features
    disparity = cv2.dilate(disparity, kernel, iterations=1)

    plot_disparity_map(disparity)

    # Remove outliers intensities from the disparity map without counting invalide disparities (min_disp - 1)
    mean = np.mean(disparity[disparity > min_disp - 1])
    std = np.std(disparity[disparity > min_disp - 1])
    # Remove outliers (values below mean - 2*std and above mean + 2*std)
    disparity[disparity < mean - 2 * std] = min_disp - 1
    disparity[disparity > mean + 2 * std] = min_disp - 1

    # Save as npy file to preserve float values
    if save:
        if not os.path.isdir("scene/disparities_map/npy"):
            os.makedirs("scene/disparities_map/npy")
        np.save(f'scene/disparities_map/npy/disparity_{id1}_{id2}.npy', disparity)
        # Save disparity map for visualization purposes
        plt.imsave(f'scene/disparities_map/disparity_{id1}_{id2}.png', disparity, cmap='jet')
    end = time.time()
    print(f"Disparity map computed in {end - start:.2f}s.")

    return disparity

def preprocess_images(image1, image2):
    # Convert to HSV
    img1_hsv = cv2.cvtColor(image1, cv2.COLOR_RGB2HSV)
    img2_hsv = cv2.cvtColor(image2, cv2.COLOR_RGB2HSV)

    # Apply CLAHE to the V channel
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    img1_hsv[:, :, 2] = clahe.apply(img1_hsv[:, :, 2])
    img2_hsv[:, :, 2] = clahe.apply(img2_hsv[:, :, 2])

    # Convert back to RGB
    img1 = cv2.cvtColor(img1_hsv, cv2.COLOR_HSV2RGB)
    img2 = cv2.cvtColor(img2_hsv, cv2.COLOR_HSV2RGB)

    return img1, img2


if __name__ == '__main__':
    from ploting import plot_points_cloud, plot_rectified_images, plot_disparity_map
    NUM_OF_IMAGES = 12
    OVERWRITE = True

    # Load images and cameras
    images = []
    cameras = []
    for i in range(NUM_OF_IMAGES):
        image = cv2.imread(f'scene/images/{i+1:02d}.jpg')
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        images.append(image)
        camera = Camera(id=i)
        camera.load(f"scene/cameras/camera{i}.txt")
        camera.image = image
        camera.image_index = '{:02d}'.format(i+1)
        cameras.append(camera)

    id1, id2 = (5, 9)
    # plot_rectified_images(cameras[id1].image, cameras[id2].image)

    print(f"Processing images {id1} and {id2}...")
    R12 = cameras[id2].R @ cameras[id1].R.T
    t12 = -cameras[id2].R @ cameras[id1].R.T @ cameras[id1].t + cameras[id2].t
    K = cameras[id1].K

    # Apply image preprocessing
    img1, img2 = preprocess_images(cameras[id1].image, cameras[id2].image)

    [H1, H2, img1Rect, img2Rect] = rectify.rectify(img1, img2, K, R12, t12)

    h = min(img1Rect.shape[0], img2Rect.shape[0])
    w = min(img1Rect.shape[1], img2Rect.shape[1])
    img1Rect, img2Rect = img1Rect[:h, :w], img2Rect[:h, :w]
    # plot_rectified_images(img1Rect, img2Rect)

    # Compute disparity map
    min_disp, max_disp = -1000, 500
    if OVERWRITE or not os.path.isfile(f'scene/disparities_map/npy/disparity_{id1}_{id2}.npy'):
        print('Computing disparity map...')
        disparity = compute_disparity_map(img1Rect, img2Rect, min_disp, max_disp)
    else:
        print('Disparity map already computed.')
        # Load disparity map from file
        disparity = np.load(f'scene/disparities_map/npy/disparity_{id1}_{id2}.npy')

    plot_disparity_map(disparity)