# Disparity Maps and Point Clouds

#import the needed libraries
import numpy as np
import matplotlib.pyplot as plt
import cv2
import matplotlib.image as isave

#read the images as grayscale
img_0 = cv2.imread("im0.png", 0)
img_1 = cv2.imread("im1.png", 0)

#convert to rgb colorspace for matplotlib
img_0_rgb = cv2.cvtColor(img_0, cv2.COLOR_BGR2RGB)
img_1_rgb = cv2.cvtColor(img_1, cv2.COLOR_BGR2RGB )

#plot the figures
plt.subplot(1,2,1)
plt.title("image 1")
plt.axis("off")
plt.imshow(img_0_rgb)

plt.subplot(1,2,2)
plt.title("image 1")
plt.axis("off")
plt.imshow(img_1_rgb)

plt.tight_layout()
plt.show()


def check_Image_Dimensions(imgA, ImgB):

    ''' Compare image dimensions by checking image shape'''
    if imgA.shape != ImgB.shape:
        raise ValueError(f"Mismtach got A {imgA.shape} vs B {ImgB.shape} ")
    else:
        print("Images match!!, Proceed")
        
#have varying number of disparities
numDisparities = [16, 32, 64, 128, 256, 512]
blockSize = [5, 7, 9, 11, 13, 15, 17, 19]

def block_matching(imgA, imgB, numDisp, bsize):
    my_stereo = cv2.StereoBM_create(numDisparities=numDisp, blockSize=bsize)
    disparity_map_raw = my_stereo.compute(imgA, imgB)

    # Inavlid disparities
    disparity_map_raw[disparity_map_raw < 0] = 0

    
    disparity_map = disparity_map_raw.astype(np.float32) / 16.0
    
    return disparity_map

#create a directory to store disparity maps for later comparison
my_maps = []

#create disparity maps
for num in numDisparities:
    for b in blockSize:

        my_disparity_map = block_matching(img_0, img_1, num, b)
        #my_maps.append((num, b, my_disparity_map))
  
        isave.imsave(f"disparity_map{num}_{b}.png", my_disparity_map, cmap="gray")
        #cv2 save image
     

  #reading pfm files for gt image
def read_pfm(file_path):
    with open(file_path, 'rb') as file:
        header = file.readline().decode('utf-8').rstrip()
        color = header == 'PF'
        
        dim_line = file.readline().decode('utf-8')
        while dim_line.startswith('#'):  
            dim_line = file.readline().decode('utf-8')
        width, height = map(int, dim_line.strip().split())
        
        scale = float(file.readline().decode('utf-8').strip())
        endian = '<' if scale < 0 else '>'
        
        data = np.fromfile(file, endian + 'f')
        shape = (height, width, 3) if color else (height, width)
        data = np.reshape(data, shape)
        data = np.flipud(data)
        
        return data      

disparity_gt = read_pfm('disp0.pfm')



 #clean disparity
disparity_gt = np.asanyarray(disparity_gt)

# Replace inf values with 0
disparity_gt = np.nan_to_num(disparity_gt, nan=0.0, posinf=0.0, neginf=0.0)

print(disparity_gt.max())
print(disparity_gt.min())
   
  


plt.imshow(disparity_gt, cmap='gray')
plt.colorbar()
plt.title('Disparity Map GT')
plt.show()
   
    


import os
import pandas as pd
import glob #pattern matching



def compute_metrics(img_0, img_1):
    # Convert to float32
    img_0 = np.array(img_0).astype(np.float32)
    img_1 = np.array(img_1).astype(np.float32)

    # Compute MSE
    mse = np.mean((img_0 - img_1) ** 2)

    if mse == 0:
        psnr = 100  
    else:
        psnr = 20 * np.log10(255.0) - 10 * np.log10(mse)

    # Compute Mean Absolute Error
    mae = np.mean(np.abs(img_0 - img_1))
   

    return psnr, mse, mae 

metrics_dict = {}
for img in glob.glob(r"C:\Users\lang_es\Desktop\CompVision\SGBM\DepthKitti\disparity*.png"):
    img_name = cv2.imread(img, 0)
    #split basename from directory and store without extension
    filename = os.path.splitext(os.path.basename(img))[0]
    #get psnr, mse, mae
    psnr, mse, mae = compute_metrics(disparity_gt, img_name)
    #print results
    #print(f"Filename {filename} has psnr of {psnr:.4f} mse {mse:.4f} mad {mae:.4f}")

    metrics_dict[img] = {
        'psnr' : psnr,
        'mse' : mse,
        'mae' : mae,
    }
    

    psnr_dict = {}
for img in glob.glob(r".\disparity*.png"):
    #load image
    my_img = cv2.imread(img, 0)
    #separate file name 
    my_file = os.path.splitext(os.path.basename(img))[0]

    psnr, mae, mse = compute_metrics(disparity_gt, my_img)

    psnr_dict[my_file] = psnr

    
#get the maximum value in the dictionary best psnr
print(max(psnr_dict, key=psnr_dict.get))

#load image
best_disp_map = cv2.imread("disparity_map256_19.png", 0)

#plot the disparity map for comparison with the gt 
plt.subplot(1,2,1)
plt.title("Best Disparity Map")
plt.axis("off")
plt.imshow(best_disp_map, cmap="plasma")

plt.subplot(1,2,2)
plt.title("Ground Truth")
plt.axis("off")
plt.imshow(disparity_gt, cmap="plasma")

plt.tight_layout()
plt.show()


#filtering remove noise and pepper while preserving the edges unlike mean/average
disparity_filtered = cv2.medianBlur(best_disp_map, 7)

plt.subplot(1,2,1)
plt.title("Best Disparity Map Without Filtering")
plt.axis("off")
plt.imshow(best_disp_map, cmap="jet")

plt.subplot(1,2,2)
plt.title("Best Disparity Map after Filtering ")
plt.axis("off")
plt.imshow(disparity_filtered, cmap="jet")

plt.tight_layout()
plt.show()


Given Camera Values

cam0=[5299.313 0 1263.818; 0 5299.313 977.763; 0 0 1]
cam1=[5299.313 0 1438.004; 0 5299.313 977.763; 0 0 1]
doffs=174.186
baseline=177.288
width=2988
height=2008
ndisp=180
isint=0
vmin=54
vmax=147
dyavg=0
dymax=0

# disparity to depth computation
def disparity_to_depth(dispariti, focal_length, baseline, doffset = 0):

    #disparity with offset
    disparity_offset = dispariti + doffset
    #define depth matrix
    depth = np.zeros_like(dispariti, dtype=np.float32)

    #check that these values are not zero to avoid division by zero 
    for i in range(dispariti.shape[0]):
        for j in range(dispariti.shape[1]):
            #check d
            d = disparity_offset[i,j]

            if d == 0:
                #set it to small value
                d = 1e-7

            #compute depth
            depth[i,j] = focal_length * baseline /(d)

    return depth

depth_map = disparity_to_depth(disparity_filtered, 5299.313, 177.288, 174.186)

plt.imshow(depth_map)


plt.subplot(1,2,1)
plt.title("Depth Map")
plt.axis("off")
plt.imshow(depth_map, cmap="jet")

plt.subplot(1,2,2)
plt.title("Disparity Map after Filtering ")
plt.axis("off")
plt.imshow(disparity_filtered, cmap="jet")

plt.tight_layout()
plt.show()


fx = 5299.313
cx = 1263.313

fy = 5299.13
cy = 977.63

#generating 3D point cloud
points = []

for i in range(depth_map.shape[0]):
    for j in range(depth_map.shape[1]):

        Z = depth_map[i,j]
        #compute x and y
        X = (j - cx) * Z / fx  
        Y = (i - cy) * Z / fy  

        points.append([X,Y,Z])

points = np.array(points)

#plot the point cloud
import open3d as o3d

#create point cloud object
point_cloud = o3d.geometry.PointCloud()
point_cloud.points = o3d.utility.Vector3dVector(points)

#visualize
o3d.visualization.draw_geometries([point_cloud])