# Copyright Oliver Zendel, AIT Austrian Institute of Technology Gmbh 2019
# Example visualization tool to illustrate interpretation of RailSem19 data usage

import sys, os, glob,  math, argparse
import numpy as np
import cv2, json
import matplotlib.pyplot as plt

rs19_label2bgr = {"buffer-stop": (70,70,70),
                "crossing": (128,64,128),
                "guard-rail": (0,255,0),
                "train-car" :  (100,80,0),
                "platform" : (232,35,244),
                "rail": (255,255,0),
                "switch-indicator": (127,255,0),
                "switch-left": (255,255,0),
                "switch-right": (127,127,0),
                "switch-unknown": (191,191,0),
                "switch-static": (0,255,127),
                "track-sign-front" : (0,220,220),
                "track-signal-front" : (30,170,250),
                "track-signal-back" : (0,85,125),
                #rail occluders
                "person-group" : (60,20,220),
                "car" : (142,0,0),
                "fence" : (153,153,190),
                "person" : (60,20,220),
                "pole" : (153,153,153),
                "rail-occluder" : (255,255,255),
                "truck" : (70,0,0)
                }

def files_in_subdirs(start_dir, pattern = ["*.png","*.jpg","*.jpeg"]):
    files = []
    for p in pattern:
        for dir,_,_ in os.walk(start_dir):
            files.extend(glob.glob(os.path.join(dir,p)))
    return files

def config_to_rgb(inp_path_config_json, default_col = [255,255,255]):
    lut = []
    inp_json = json.load(open(inp_path_config_json, 'r'))
    for c in range(3): #for each color channel
        lut_c =[l["color"][c] for l in inp_json["labels"]]+[default_col[c]]*(256-len(inp_json["labels"]))
        lut.append(np.asarray(lut_c, dtype=np.uint8))
    return lut

def corss_hatch_rail(im_vis, coords, color_left=(255,255,0), color_right=(127,127,0)):
    ml = min(len(coords[0]), len(coords[1]))
    for i in range(ml):
        midpnt = ((coords[0][i][0]+coords[1][i][0])//2, (coords[0][i][1]+coords[1][i][1])//2)
        cv2.line(im_vis, tuple(coords[0][i]), midpnt, color_left)
        cv2.line(im_vis, midpnt, tuple(coords[1][i]), color_right)

def json_to_img(inp_path_json, line_thickness=2):
    inp_json = json.load(open(inp_path_json, 'r'))
    im_json = np.zeros((inp_json["imgHeight"], inp_json["imgWidth"],3), dtype=np.uint8)
    for obj in inp_json["objects"]:
        col = rs19_label2bgr.get(obj["label"],[255,255,255])
        if "boundingbox" in obj:
            cv2.rectangle(im_json, tuple(obj["boundingbox"][0:2]), tuple(obj["boundingbox"][2:4]), col, line_thickness)
        elif "polygon" in obj:
            pnts_draw = np.around(np.array(obj["polygon"])).astype(np.int32)
            cv2.polylines(im_json, [pnts_draw], True, col, line_thickness)
        elif "polyline-pair" in obj:
            #left rail of a rail pair has index 0, right rail of a rail pair has index 1
            rails_draw = [np.around(np.array(obj["polyline-pair"][i])).astype(np.int32) for i in range(2)]
            corss_hatch_rail(im_json, obj["polyline-pair"],  rs19_label2bgr['switch-left'], rs19_label2bgr['switch-right'])
            cv2.polylines(im_json, rails_draw, False, col)
        elif "polyline" in obj:
            rail_draw = np.around(np.array(obj["polyline"])).astype(np.int32)
            cv2.polylines(im_json, [rail_draw], False, col, line_thickness)
    return im_json, inp_json["frame"]

def get_joined_img(inp_path_json, jpg_folder, uint8_folder, lut_bgr, blend_vals=[0.65,0.25,0.1]):
    im_json, frameid = json_to_img(inp_path_json, line_thickness=2) #visualize geometric annotations
    inp_path_jpg = os.path.join(jpg_folder,frameid+".jpg")
    inp_path_uint8 = os.path.join(uint8_folder,frameid+".png")
    im_jpg = cv2.imread(inp_path_jpg) #add intensity image as background
    im_id_map = cv2.imread(inp_path_uint8,cv2.IMREAD_GRAYSCALE) #get semantic label map
    im_id_col = np.zeros((im_id_map.shape[0], im_id_map.shape[1], 3), np.uint8)
    for c in range(3):
        im_id_col[:,:,c] = lut_bgr[c][im_id_map] #apply color coding
    
    #plt.imshow(im_id_col)
    #plt.show()

    blend_sources = (im_jpg*blend_vals[0]+im_id_col*blend_vals[1]+im_json*blend_vals[2]).astype(np.uint8)
    return (blend_sources) #blend all three data sources

def vis_all_json(json_folder, jpg_folder, uint8_folder, inp_path_config_json):
    all_json = files_in_subdirs(json_folder, pattern = ["*.json"])
    lut_bgr = config_to_rgb(inp_path_config_json, default_col = [255,255,255])[::-1] #we need to swap color channels as opencv uses BGR
    curr_idx, retKey = 0, ord('a')
    while retKey > ord(' '):
        im_vis = get_joined_img(all_json[curr_idx], jpg_folder, uint8_folder, lut_bgr)
        cv2.putText(im_vis, all_json[curr_idx],(0,25), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 2)
        cv2.imshow("RailSem19 Annotation Visualization (use 'a' or 'd' to navigate, ESC to quit)",im_vis)
        retKey = cv2.waitKey(-1) #use 'a' and 'd' to scroll through frames
        curr_idx = curr_idx - 1 if retKey == ord('a') else curr_idx + 1
        curr_idx = (curr_idx + len(all_json)) % len(all_json) #wrap around
    return 0

def main(argv=sys.argv[1:]):
    parser = argparse.ArgumentParser()
    parser.add_argument('--jpgs', type=str, default="RailNet_DT/rs19_val/jpgs/rs19_val",
                        help="Folder containing all RGB intensity images")
    parser.add_argument('--jsons', type=str, default="RailNet_DT/rs19_val/jsons/rs19_val",
                        help="Folder containing all geometry-based annotation files")
    parser.add_argument('--uint8', type=str, default="RailNet_DT/rs19_val/uint8/rs19_val",
                        help="Folder containing all dense semantic segm. id maps")
    parser.add_argument('--config', type=str, default="RailNet_DT/rs19_val/rs19-config.json",
                        help="Path to config json for dense label map interpretation")
    args = parser.parse_args(argv)
    return vis_all_json(args.jsons, args.jpgs, args.uint8, args.config)
    
if __name__ == "__main__":
    sys.exit(main())