utils.py

import cv2
import numpy as np
from point_cloud_generator import PointCloudGenerator

# pcd_generator = PointCloudGenerator()

def resize(image):
    """
    resize the input nd array
    """
    h, w = image.shape[:2]
    if h > w:
        return cv2.resize(image, (480, 640)) 
    else:
        return cv2.resize(image, (640, 480)) 

def get_masked_depth(depth_map, mask):
    masked_depth_map = depth_map*mask
    pixel_depth_vals = masked_depth_map[masked_depth_map>0]
    mean_depth = np.mean(pixel_depth_vals)
    return masked_depth_map, 1-mean_depth

def draw_depth_info(image, depth_map, objects_data):
    image = image.copy()
    # object data -> [cls_id, cls_name, cls_center, cls_mask, cls_clr]
    for data in objects_data:
        center = data[2]
        mask = data[3]
        _, depth = get_masked_depth(depth_map, mask)
        cv2.rectangle(image, (center[0]-15, center[1]-15), (center[0]+(len(str(round(depth*10, 2))+'m')*12), center[1]+15), data[4], -1)
        cv2.putText(image, str(round(depth*10, 2))+'m', (center[0]-5, center[1]+5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
    
    return image

def generate_obj_pcd(depth_map, objects_data):
    objs_pcd = []
    pcd_generator = PointCloudGenerator()

    for data in objects_data:
        mask = data[3]
        cls_clr = data[4]
        masked_depth = depth_map*mask
        # generating point cloud using masked depth
        pcd = pcd_generator.generate_point_cloud(masked_depth)
        objs_pcd.append((pcd, cls_clr))
    return objs_pcd