obj_avoidance.py

"""
implement an obstacle avoidance algorithm

hyperparameters:
1. central part: [600:680], i have tested many ranges, this one performs best; it should be changed if you change the size of images
2. depth threshold: i set 2000
3. threshold of the number of pixels: i set 500
4: object_classes = [2,3,4,5,6]
5. label_dict = { 1:'floor',2: 'furniture', 3:'objects',  4: 'person', 5: 'wall', 6:'door', 7:'ceiling'}

note:
we should make sure the size of pred and d1 are the same
"""
import os
import sys
import time
import copy
import re

import matplotlib.pyplot as plt
import cv2
from PIL import Image
import numpy as np
from skimage import measure

# by maggir
#label_dict = { 1: 'floor', 2: 'furniture', 3:'objects',  4: 'person', 5: 'wall', 6:'door', 7:'ceiling'}  
# note that colors[key] --> names[key+1] (by me)
#label_dict = {1: 'wall', 2: 'floor', 3: 'plant', 4: 'ceiling', 5: 'furniture', 6: 'person', 7: 'door', 8: 'objects'}

def run_avoidance(d1_img, seg_idx, depth_threshold = 500, visible_width = 60):
    """
    input:
        d1_img -- np array, 1 channel depth (uint16)
        seg_idx -- np array, segmentation output with index only
        depth_threshold -- int, for apply masking by depth, *for depth in uint16
        visible_width -- 90, number of pixel up front you wanna be visible
        ** make sure d1_img, seg_idx have same size
    output:
        obj_tup -- tuple, (class idx, distance, (x, y))
        obj_img -- np array -- image showing single object only
    """
    assert d1_img.shape == seg_idx.shape, '[ERROR] depth image shape not align with segmentation index'
    # omit floor (idx = 1, name = 2) and ceiling (idx = 3, name = 4)
    seg_idx[seg_idx == 1] = 0
    seg_idx[seg_idx == 3] = 0
    # find connected components (index per instance)
    inst_idx = measure.label(seg_idx, connectivity = 2)
    # get closest object info
    obj_tup  = get_obj_info(d1_img, seg_idx, inst_idx, depth_threshold = depth_threshold, visible_width = visible_width)
    # get the closest object image
    obj_img = get_obj_img(d1_img, inst_idx, obj_tup[0])
    return obj_tup, obj_img


def get_obj_info(d1_img, seg_idx, inst_idx, depth_threshold, visible_width):
    """
    input:
        d1_img -- np array, 1 channel depth (uint8)
        seg_idx -- np array, segmentation index map
        inst_idx -- np array, output from connected components
        depth_threshold -- int, for apply masking by depth
        visible_width -- 90, number of pixel up front you wanna be visible
        ** make sure d1_img, seg_idx have same size
    output:
        (min_inst_idx, min_cls_idx, min_dist) -- (instance index, class index, its distance)
    """
    _, w = seg_idx.shape
    # set up front angle width
    lower_limit = int(w/2 - visible_width/2)
    upper_limit = int(w/2 + visible_width/2)
    # apply filter on distance and angle
    filter_inst_idx = (d1_img < depth_threshold) * inst_idx
    for i in np.unique(filter_inst_idx):
        if i != 0:
            if len(np.where(filter_inst_idx != 0)[0]) < 850: filter_inst_idx[np.where(filter_inst_idx != 0)] = 0
    # find the closest distance
    filter_inst_idx[:, :lower_limit] = 0
    filter_inst_idx[:, upper_limit:] = 0
    # find closeset distance
    min_inst_idx, min_cls_idx, min_dist = None, None, float('inf')
    for idx in np.unique(filter_inst_idx):
        idx_locs = np.where(inst_idx == idx)
        if len(idx_locs[0]) == 0:
            continue
        loc = (idx_locs[0][0], idx_locs[1][0])
        cls_idx = seg_idx[loc[0], loc[1]]
        # skip 0 index (background)
        if cls_idx == 0:
            continue
        depth = d1_img[idx_locs]
        # remove noise in depth map
        depth[depth == 0] = depth.max()
        depth[depth < 150] = depth.max()
        #dist = np.sort(depth, axis = None)

        dist = np.sort(depth, axis = None)[:20].mean()
        if dist < min_dist:
            min_inst_idx = idx
            min_cls_idx = cls_idx
            min_dist = dist
    return min_inst_idx, min_cls_idx, min_dist


def get_obj_img(d1_img, inst_idx, target_idx):
    """
    visualise the whole closest object 

    input:
        d1_img -- np array, 1 channel depth
        seg_idx -- np array, segmentation output with index only
        target_idx -- int, target index of the instance
    output:
        obj_img -- np array, showing closest object only (having same dim as input)
    """
    if target_idx is None:
        return np.zeros(shape = d1_img.shape, dtype = np.uint8)
    else:
        # sharpen the intensity by * 30
        return (inst_idx == target_idx) * d1_img * 30