Skip to content
Snippets Groups Projects
Select Git revision
  • master
1 result

714.nfa

Blame
  • Code owners
    Assign users and groups as approvers for specific file changes. Learn more.
    check_depth.py 2.44 KiB
    """
    Do some checks for the rendered depth images. 
    (total number of objects and poses for each object, and if they are 16-bit png file, and if "non-all-background")
    """
    import os 
    os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
    import numpy as np
    import cv2 
    import tqdm
    import glob
    from pathlib import Path
    
    # return True if the image is one-channel
    def isOneChannel(img):
        return True if len(img.shape) == 2 else False
    
    # return True if the image contains an object (not only background)
    def isNotEmpty(img):
        values_count = len(set(img.ravel().tolist()))
        return True if values_count > 20 else False 
        # the threshold should be 1, but let's set it higher so it can detect some weird objects in the dataset
    
    def isNormailized(img, d_min = 0.1):
        img_max = img.max()
        img_min = img.min()
        return True if img_min >= 0 and img_max <= 65535 else False 
    
    def print_path(path):
        print(path.split("/")[-4], path.split("/")[-3], path.split("/")[-1])
    
    # loop through all the .png files under start_dir 
    def find_png_files(start_dir):
        # The '**' pattern means "this directory and all subdirectories"
        # The 'recursive=True' makes sure the pattern '**' is applied recursively
        for img_path in tqdm.tqdm(glob.glob(f"{start_dir}/**/*.png", recursive=True)):
            img = cv2.imread(img_path,  cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) 
            checks = [isOneChannel, isNotEmpty, isNormailized]
            for check in checks:
                if not check(img):
                    print_path(img_path)
                    print(f"| {check.__name__}: False")
    
    
    # check if the number of rendered depth images matches its original dataset (e.g. SRN car_train)
    def check_obj_num(start_dir, obj_num, pose_num):
        p = Path(start_dir)
        # check if number of objects are correct 
        detected_obj_num = len([sub_dir for sub_dir in p.iterdir() if sub_dir.is_dir()])
        if detected_obj_num != obj_num:
            raise ValueError(f"Expected {obj_num} objects, but found {detected_obj_num} objects.")
        # check of number of images (poses) are correct for an object 
        for sub_dir in p.iterdir():
            depth_dir = sub_dir / "depth"
            detected_pose_num = len(list(depth_dir.glob('*.png')))
            if detected_pose_num != pose_num:
                raise ValueError(f"Expected {pose_num} png files, but found {detected_pose_num} files. \n {depth_dir}")
    
    start_dir = "/nodes/astra/fastwork/wang_c/SRN_depth/cars_test"
    check_obj_num(start_dir, obj_num = 704, pose_num = 251)
    find_png_files(start_dir)