Pipeline for fetching frames from Intel's RealSense using PyRealSense

dataDump.py

#!/usr/bin/env python

"""
This code collects the frames from Intel RealSense camera
and dumps them into the data folder.

Frames that will be fetched:
1) Color
2) Depth
3) Color aligned on depth
4) Depth aligned on color

If the directories are not present, this code will 
automatically create them.

The directory structure will be as follows:

data
├── color
├── depth
├── cad
└── dac
"""

import os
import sys
import cv2
import pyrealsense as pyrs


def create_directory(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)
        return 1
    else:
        print(directory, '- already exists')
        return 0
    

def check_directories():
    try:
        create_directory('./data')
        create_directory('./data/depth')
        create_directory('./data/color')
        create_directory('./data/cad')
        create_directory('./data/dac')
    except:
        print("Unexpected error:", sys.exc_info()[0])
        return -1
    return 0

def main():
    file_structure = check_directories()
    if file_structure == -1:
        print('\nERROR: Directories can\'t be created, error thrown')
        return -1
    else:
        print('\nDirectories created successfully...\nLaunching camera module...')
    
    # Fire camera & launch streams
    pyrs.start()
    cam = pyrs.Device(device_id = 0, streams = [pyrs.stream.ColorStream(fps=60),
                                                pyrs.stream.DepthStream(fps=60),
                                                pyrs.stream.CADStream(fps=60),
                                                pyrs.stream.DACStream(fps=60)])
    scale = cam.depth_scale * 1000

    # Some important variables
    flag_save_frames = False
    file_num = 0

    # Start fetching Buffer
    print('Starting Buffer...')
    while(True):
        cam.wait_for_frames()
        current_color = cam.color[:,:,::-1]
        current_depth = cam.depth * scale
        current_cad = cam.cad[:,:,::-1]
        current_dac = cam.dac * scale
        
        cv2.imshow('Color',current_color)
        cv2.imshow('Depth',current_depth/1000)
        cv2.imshow('CAD',current_cad)
        cv2.imshow('DAC',current_dac/1000)

        

        if flag_save_frames:
            num = format(file_num, '08')
            cv2.imwrite('./data/depth/' + str(num) + '.png', cam.depth)
            cv2.imwrite('./data/color/' + str(num) + '.png', current_color)
            cv2.imwrite('./data/dac/' + str(num) + '.png', cam.dac)
            cv2.imwrite('./data/cad/' + str(num) + '.png', current_cad)
            file_num += 1

        k = cv2.waitKey(1)
        if k == ord('q'):
            print('Q Pressed...\nEnding execution')
            break
        if k == ord('f'):
            if flag_save_frames:
                print('F Pressed...\nStopped fetching frames...')
                flag_save_frames = False
            else:
                print('F Pressed...\nStarted fetching frames...')
                flag_save_frames = True

    cam.stop()
    pyrs.stop()

    return 0


if __name__ == '__main__':
    print(__doc__)
    main()