bag2dataset.py

from typing import Type

import rosbag
import argparse
import multiprocessing
import gc
# from multiprocessing import cpu_count, Pool
# from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

from tqdm import tqdm
import numpy as np
import numba as nb
from colour_demosaicing import demosaicing_CFA_Bayer_Menon2007
from colour import Lab_to_XYZ, XYZ_to_Lab
from sensor_msgs.msg import Image, CompressedImage
# from sklearn.preprocessing import PolynomialFeatures
# from sklearn.linear_model import LinearRegression
# from sklearn.pipeline import Pipeline
# from rospy_message_converter import json_message_converter
import scipy.optimize as opt
from libtiff import TIFF
# import fcntl
import os
import io
# import ffmpeg
import cv2
# import png
from time import perf_counter


def whitebalance(img):
    pixels = img.reshape((img.shape[0] * img.shape[1], -1)).T
    transformed = wb_transform @ pixels
    return np.clip(transformed.T.reshape(img.shape), 0., 1.)

def autowhitebalance(img, config=None, perc=0.6):
    # equivalent to GIMP auto white balance
    # exclude top/bottom 0.05% of pixels in each channel and stretch the rest to fit 0-1
    # GIMP docs say 0.05% but it's actually 0.6%... fml...
    if perc > 50:
        perc = 100 - perc
    if config is not None:
        lower, upper = config
    else:
        lower, upper = [None] * 3, [None] * 3
    pixels = img.reshape((img.shape[0] * img.shape[1], 3))
    out = np.empty_like(img)
    for i in range(3):
        if lower[i] is None:
            lower[i] = np.percentile(pixels[:, i],  perc)
        if upper[i] is None:
            upper[i] = np.percentile(pixels[:, i], 100. - perc)
        out[:, :, i] = np.clip(img[:, :, i], lower[i], upper[i])
        out[:, :, i] = (out[:, :, i] - lower[i]) / (upper[i] - lower[i])
    config = lower, upper
    return out, config


def linear_to_sRGB(img):
    out = np.empty_like(img)
    mask = img <= 0.0031308
    out[mask] = img[mask] * 12.92
    out[~mask] = np.power(img[~mask], 1/2.4) * 1.055 - 0.055
    # mask = img > 0
    # out[~mask] = 0.
    # out[mask] = np.power(img[mask], 1./2.2)
    return out


base_link_to_downward_rgb_frame = {
    "header": {
        "seq": 0,
        'stamp': {
            "secs": 0,
            "nsecs": 0
        },
        "frame_id": '/warpauv_2/base_link'
    },
    'child_frame_id': '/warpauv_2/cameras/downward/rgb',
    'transform': {
        'translation': {
            'x': 0.1125,
            'y': 0.0,
            'z': -0.195
        },
        'rotation': {
            'x': -0.7071068,
            'y': 0.7071068,
            'z': 0.0,
            'w': 0.0
        }
    }
}

base_link_to_forward_rgb_frame = {
    "header": {
        "seq": 0,
        'stamp': {
            "secs": 0,
            "nsecs": 0
        },
        "frame_id": '/warpauv_2/base_link'
    },
    'child_frame_id': '/warpauv_2/cameras/forward/rgb',
    'transform': {
        'translation': {
            'x': 0.179,
            'y': 0.0,
            'z': -0.165
        },
        'rotation': {
            'x': -0.6532815,
            'y': 0.6532815,
            'z': -0.2705981,
            'w': 0.2705981
        }
    }
}

''' Packing scheme for RAW10 - MIPI CSI-2
- 4 pixels: p0[9:0], p1[9:0], p2[9:0], p3[9:0]
- stored on 5 bytes (byte0..4) as:
| byte0[7:0] | byte1[7:0] | byte2[7:0] | byte3[7:0] |          byte4[7:0]             |
|    p0[9:2] |    p1[9:2] |    p2[9:2] |    p3[9:2] | p3[1:0],p2[1:0],p1[1:0],p0[1:0] |
'''
# Optimized with 'numba' as otherwise would be extremely slow (55 seconds per frame!)
@nb.njit(nb.uint16[::1] (nb.types.Array(nb.types.uint8, 1, 'C', readonly=True),
                         nb.uint16[::1], nb.boolean), parallel=True, cache=True)
def unpack_raw10(input, out, expand16bit):
    lShift = 6 if expand16bit else 0

   #for i in np.arange(input.size // 5): # around 25ms per frame (with numba)
    for i in nb.prange(input.size // 5): # around  5ms per frame
        b4 = input[i * 5 + 4]
        out[i * 4]     = ((input[i * 5]     << 2) | ( b4       & 0x3)) << lShift
        out[i * 4 + 1] = ((input[i * 5 + 1] << 2) | ((b4 >> 2) & 0x3)) << lShift
        out[i * 4 + 2] = ((input[i * 5 + 2] << 2) | ((b4 >> 4) & 0x3)) << lShift
        out[i * 4 + 3] = ((input[i * 5 + 3] << 2) |  (b4 >> 6)       ) << lShift

    return out


# def extract_h264_topic(bag: rosbag.Bag, image_topic: str, output_dir: str):
#     """Extract a h264 file from a rosbag.
#     """
#     topic: str
#     msg: CompressedImage
#     process = (
#         ffmpeg
#         .input('pipe:', f='h264', framerate='6')
#         .video
#         .output('pipe:', format='rawvideo', pix_fmt='bgr24')
#         .overwrite_output()
#         .run_async(pipe_stdin=True, pipe_stdout=True)
#     )
#     height, width = 2160, 3840
#
#     fd = process.stdout.fileno()
#     fl = fcntl.fcntl(fd, fcntl.F_GETFL)
#     fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
#
#     for topic, msg, t in tqdm(bag.read_messages(topics=[image_topic, '/tf', '/tf_static']), desc="Reading bag"):
#         if topic != image_topic:
#             continue
#         # stream = io.BytesIO(msg.data)
#         process.stdin.write(msg.data)
#         # process.stdin.flush()
#         # in_bytes, errs = process.communicate(msg.data, timeout=1)
#
#         in_bytes = process.stdout.read(width * height * 3)
#         while len(in_bytes) >= width * height * 3:
#             in_frame = np.frombuffer(in_bytes[:width*height*3], np.uint8).reshape([height, width, 3])
#             in_bytes = in_bytes[width*height*3:]
#             cv2.imshow('frames', in_frame)
#             if cv2.waitKey(100) & 0xFF == ord('q'):
#                 break
#         else:
#             print("No more data to read")
#     process.stdin.close()
#     process.wait()
#     cv2.destroyAllWindows()


def process_packed_bayer_img(packed, size, color_transform, output_path, debug=False, wb_config=None):
    # print(f"Processing {packed.size} bytes, saving to {output_path}")
    width, height = size
    bayer_img = np.empty(packed.size * 4 // 5, dtype=np.uint16)
    bayer_img: np.ndarray = unpack_raw10(packed, bayer_img, expand16bit=True).reshape((height, width))
    color_img: np.ndarray = demosaicing_CFA_Bayer_Menon2007(bayer_img.astype(np.float32) / np.iinfo(np.uint16).max,
                                                            "RGGB")
    corrected_img: np.ndarray = (
                color_transform @ np.concatenate((color_img.reshape((width * height, 3)).T, np.ones((1, width * height))),
                                   axis=0)).T.reshape(color_img.shape)
    wb_img, wb_config = autowhitebalance(corrected_img, config=wb_config, perc=0.8)
    gamma_img = linear_to_sRGB(wb_img)
    if debug and output_path is not None:
        base_path, ext = os.path.splitext(output_path)
        debug_fname = lambda label: base_path + "-" + label + ext
        bayer_outfile = TIFF.open(debug_fname("1bayer"), mode='w')
        bayer_outfile.write_image(bayer_img, compression='deflate')
        debayered_outfile = TIFF.open(debug_fname("2debayered"), mode='w')
        debayered_outfile.write_image(color_img, write_rgb=True)
        corrected_outfile = TIFF.open(debug_fname("3corrected"), mode='w')
        corrected_outfile.write_image(corrected_img, write_rgb=True)
        wb_outfile = TIFF.open(debug_fname("4whitebalanced"), mode='w')
        wb_outfile.write_image(wb_img, write_rgb=True)
        gamma_outfile = TIFF.open(debug_fname("5gammacorrected"))
        gamma_outfile.write_image(gamma_img, write_rgb=True)
    del packed, bayer_img, color_img, corrected_img, wb_img
    # print("Writing output file to", output_path)
    if output_path is not None:
        out_img = cv2.addWeighted(gamma_img, alpha=np.iinfo(np.uint8).max, src2=0, beta=0, gamma=0, dtype=cv2.CV_8U)
        corrected_outfile = TIFF.open(output_path, mode='w')
        corrected_outfile.write_image(out_img, compression='deflate', write_rgb=True)
        del gamma_img, out_img
    gc.collect()
    return wb_config


def bag_iterator(progress, bag, topic, output_dir, size, color_transform, debug=False, t_start=0):
    skipped = 0
    wb_config = None
    for topic, msg, t in bag.read_messages(topics=[topic]):
        if t.to_sec() < t_start:
            skipped += 1
            continue
        packed = np.frombuffer(msg.data, dtype=np.uint8)
        if wb_config is None:
            wb_config = process_packed_bayer_img(packed, size, color_transform, output_path=None, debug=debug)
        progress.total += 1
        progress.update(0)
        yield msg.header, packed, t, output_dir, size, color_transform, wb_config, debug
    print(f"Skipped {skipped} messages")


def process_msg(args):
    header, packed, t, output_dir, size, color_transform, wb_config, debug = args
    corrected_fpath = os.path.join(output_dir, f"{t.secs:09}.{t.nsecs:09}.tiff")
    process_packed_bayer_img(packed, size, color_transform, output_path=corrected_fpath, debug=debug, wb_config=wb_config)
    return header, corrected_fpath

def extract_raw_topic(bag, image_topic, output_dir, color_transform, t_start=0, max_cpus=multiprocessing.cpu_count(), debug=False):
    raw_images = []
    topic: str
    # msg: Image
    size = 3840, 2160  # width, height

    msg: Image
    with multiprocessing.get_context('forkserver').Pool(max_cpus) as pool:
        progress = tqdm(desc="Extracting and processing raw images", total=0)
        for result in pool.imap_unordered(process_msg, bag_iterator(progress, bag, image_topic, output_dir, size, color_transform, debug=debug, t_start=t_start), chunksize=2):
            raw_images.append(result)
            progress.update(1)
        progress.close()

        gc.collect()
        pool.close()
        pool.join()
    return raw_images

if __name__ == "__main__":


    parser = argparse.ArgumentParser(description="Extract images from a ROS bag.")
    parser.add_argument("odometry_bag", help="Input odometry ROS bag.")
    parser.add_argument("image_bag", help="Input image ROS bag.")
    parser.add_argument("output_dir", help="Output directory.")
    parser.add_argument("image_topic", help="Image topic")
    parser.add_argument("--t_start", default=0., type=float, help="Ignore messages before this time")
    parser.add_argument("--threads", default=multiprocessing.cpu_count(), type=int)
    parser.add_argument("--debug", default=False, action="store_true")
    parser.add_argument("--color-correct", default=True, action="store_true")
    parser.add_argument("--no-color-correct", dest='color_correct', action="store_false")
    # parser.add_argument("odometry_topic", help="Image topic")

    args = parser.parse_args()

    if args.color_correct:
        reference_LAB = np.array([
            [97, 0, 1],
            [73, 0, 0],
            [62, 0, 0],
            [50, 0, 0],
            [38, 0, 0],
            [23, 0, 0],
            [48, 59, 39],
            [92, 1, 95],
            [64, -40, 54],
            [57, -41, -42],
            [18, -3, -25],
            [49, 60, -3],
            [41, 51, 26],
            [61, 29, 57],
            [52, -24, -24],
            [52, 47, -14],
            [69, 14, 17],
            [64, 12, 17]
        ], dtype=float)

        patches_LAB = np.array([  # warpauv_2_xavier3_RAW_2022-11-02-10-25-39.bag @ 2725.797333449s
            [91, -48, -14.3],
            [76.8, -52.8, 11.2],
            [58.2, -41.3, 8.7],
            [42.4, -31.2, 6.2],
            [32.8, -24.5, 4.3],
            [22.4, -16.8, 1.7],
            [36.4, -27.7, 9.6],
            [88.6, -70.7, 48.7],
            [72.1, -58.4, 36.1],
            [76.7, -44.4, -8.8],
            [26.7, -18.7, -0.4],
            [38.5, -19.3, -13.1],
            [29.2, -22.1, 5.1],
            [41.4, -33.4, 15.7],
            [53.9, -36.3, 1.5],
            [40.5, -23.1, -8.9],
            [57.1, -42.3, 13.6],
            [50.6, -38.2, 12.8]
        ], dtype=float)


        def lsq_color_correct(obs, ref):
            aug_obs = np.concatenate((np.ones((obs.shape[0], 1)), obs), axis=1)
            delta = np.linalg.inv(aug_obs.T @ aug_obs)
            A = delta @ aug_obs.T @ ref
            return A


        # def poly_color_correct(obs, ref, order=4):
        #     model = Pipeline([('poly', PolynomialFeatures(degree=order)),
        #                       ('linear', LinearRegression(fit_intercept=False))])
        #     params = []
        #     for c in range(3):
        #         model.fit(obs, ref[:, c])
        #         params.append(model.named_steps['linear'].coef_)
        #     return np.stack(params, axis=0)

        linear_sRGB_to_XYZ = np.array([[0.4124564, 0.3575761, 0.1804375],
                                       [0.2126729, 0.7151522, 0.0721750],
                                       [0.0193339, 0.1191920, 0.9503041]])
        B = linear_sRGB_to_XYZ
        Bext = np.eye(4)
        Bext[0:3, 0:3] = B
        Binv = np.linalg.inv(B)

        Y = Lab_to_XYZ(reference_LAB).T
        X = np.concatenate((Lab_to_XYZ(patches_LAB), np.ones((18, 1))), axis=1).T
        A = np.zeros((3, 4))
        Ymean = Y.mean(axis=0)
        Xmean = X.mean(axis=0)
        A[0, 0] = Ymean[0] / Xmean[0]
        A[1, 1] = Ymean[1] / Xmean[1]
        A[2, 2] = Ymean[2] / Xmean[2]

        obj = lambda a: np.linalg.norm(reference_LAB - XYZ_to_Lab(X.T @ a.reshape((3, 4)).T), axis=1).flatten()
        print(f"Initial: {obj(A.flatten())}")
        res = opt.least_squares(obj, A.flatten(), loss='linear', jac='3-point')
        Aopt = res.x.reshape((3, 4))
        print(f"Final: {obj(Aopt.flatten())}")

        Aopt_new = np.roll(lsq_color_correct(obs=Lab_to_XYZ(patches_LAB), ref=Lab_to_XYZ(reference_LAB)).T, -1, axis=1)
        # Aopt_poly = poly_color_correct(obs=Lab_to_XYZ(patches_LAB), ref=Lab_to_XYZ(reference_LAB))
        T = Binv @ Aopt @ Bext
        T_new = Binv @ Aopt_new @ Bext

        gray18_xyz_ref = Y[:, 3]
        gray18_xyz = X[:, 3]
        gray18_rgb = Binv @ Aopt @ gray18_xyz
        gray18_rgb_ref = Binv @ gray18_xyz_ref
        wb_transform = np.diag(gray18_rgb_ref / gray18_rgb)
    else:
        T = np.eye(4)[:3, :]

    image_bag = rosbag.Bag(args.image_bag, "r")
    # if args.image_topic.endswith("/h264"):
    #     extract_h264_topic(image_bag, args.image_topic, args.output_dir)
    if args.image_topic.endswith("/image_bayer"):
        extract_raw_topic(image_bag, args.image_topic, args.output_dir, T, t_start=args.t_start, max_cpus=args.threads, debug=args.debug)
    image_bag.close()

    odometry_bag = rosbag.Bag(args.odometry_bag, "r")
    odometry_bag.close()