# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# !pip install -e .

import pickle

from PIL import Image

# +
# from FusionTransformer.data.utils.visualize import draw_points_image_labels
# -

from FusionTransformer.data.utils.visualize import SEMANTIC_KITTI_COLOR_PALETTE

import numpy as np
import cv2

import matplotlib.pyplot as plt

labels ={
  0 : "unlabeled",
  1 : "outlier",
  10: "car",
  11: "bicycle",
  13: "bus",
  15: "motorcycle",
  16: "on-rails",
  18: "truck",
  20: "other-vehicle",
  30: "person",
  31: "bicyclist",
  32: "motorcyclist",
  40: "road",
  44: "parking",
  48: "sidewalk",
  49: "other-ground",
  50: "building",
  51: "fence",
  52: "other-structure",
  60: "lane-marking",
  70: "vegetation",
  71: "trunk",
  72: "terrain",
  80: "pole",
  81: "traffic-sign",
  99: "other-object",
  252: "moving-car",
  253: "moving-bicyclist",
  254: "moving-person",
  255: "moving-motorcyclist",
  256: "moving-on-rails",
  257: "moving-bus",
  258: "moving-truck",
  259: "moving-other-vehicle"
}

with open("/home/user/SemanticKitti/preprocessed/08/scan_data_1.pkl", 'rb') as f:
    filter_data = pickle.load(f)

type(filter_data)

filter_data.keys()

filter_data["camera_path"]

img_path = "/home/user/SemanticKitti/" + filter_data["camera_path"]

image_width = 1226
image_height = 370


# +
# image = Image.open(img_path).crop((0, 0, image_width, image_height))

# +
# draw_points_image_labels(img=image, 
#                          img_indices=filter_data["points_img"], 
#                          seg_labels=filter_data["seg_labels"], 
#                          show=True, 
#                          color_palette_type='SemanticKITTI_long', 
#                          point_size=0.5)
# -

def draw_points_image_labels(img, img_indices, seg_labels, show=True, point_size=0.5, sp=None):
    img = img.copy()
    color_palette = SEMANTIC_KITTI_COLOR_PALETTE
    color_palette = np.array(color_palette) #/ 255.
    seg_labels[seg_labels == 0] = len(color_palette) - 1
    colors = color_palette[seg_labels]
    if sp is not None:
        colors= colors[seg_labels == sp]
        img_indices = img_indices[seg_labels == sp]
    
    for i in range(img_indices.shape[0]):
        x, y = img_indices[i, :].astype(np.int64)
        cv2.circle(img, center=(y,x), radius=2, color=colors[i].tolist())

    return img


img = cv2.imread(img_path)

img.shape

img = img[:image_height, :image_width, :]

img.shape

filter_data["points_img"].shape

filter_data["seg_labels"].shape

to_names = lambda l: [ (labels[i],i) for i in np.unique(l).tolist()]

to_names(filter_data["seg_labels"])

colored_img = draw_points_image_labels(img=img, 
                         img_indices=filter_data["points_img"], 
                         seg_labels=filter_data["seg_labels"], 
                         show=True,  
                         point_size=0.5,
                         sp=30)

colored_img = cv2.cvtColor(colored_img, cv2.COLOR_BGR2RGB)

fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(10, 10))
ax[1].imshow(colored_img)
ax[0].imshow(img)

StarcoderdataJupyterScriptsDedupFiltered