import cv2
import torch
import numpy as np

from depth_anything_v2.dpt import DepthAnythingV2

model_configs = {
    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
}

encoder = 'vitl' # or 'vits', 'vitb'
dataset = 'pbr' # 'hypersim' for indoor model, 'vkitti' for outdoor model
max_depth = 1 # 20 for indoor model, 80 for outdoor model

model = DepthAnythingV2(**{**model_configs[encoder], 'max_depth': max_depth})

# Load checkpoint and handle unexpected keys
checkpoint = torch.load(f'checkpoints/model2.pth', map_location='cpu')
print("Keys in checkpoint:", checkpoint.keys())

# Skip unexpected keys
expected_keys = ['model']
state_dict = {}
for key in checkpoint.keys():
    if key not in ['optimizer', 'epoch', 'previous_best']:
        state_dict = checkpoint[key]
        print(f"Using weights from key: {key}")
    else:
        print(f"Skipping unexpected key: {key}")

# Handle module prefix if present
my_state_dict = {}
for key in state_dict.keys():
    new_key = key.replace('module.', '')
    my_state_dict[new_key] = state_dict[key]

model.load_state_dict(my_state_dict)
model.eval()

raw_img = cv2.imread('image.jpg')
depth = model.infer_image(raw_img) # HxW depth map in meters in numpy

# Normalize depth for visualization (0-255)
depth_normalized = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)

# Apply colormap for better visualization
depth_colormap = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)

# Save both raw depth and colored depth
cv2.imwrite('depth_raw.png', depth_normalized)
cv2.imwrite('depth_colored.png', depth_colormap)
print("Saved depth maps as 'depth_raw.png' and 'depth_colored.png'")