{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Car pose estimation\n---\n\n- v0.1: Created functions to visualize the data"},{"metadata":{"_uuid":"89e13c4f-f1e4-4897-b861-84dc6f15b869","_cell_guid":"c65e8aae-8cd2-49c0-bb5b-a5938c153cd9","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport math\n\nfrom matplotlib import pyplot as plt\n\nimport cv2\nfrom PIL import Image\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data visualization stuff\n** Defining constants and utility functions **"},{"metadata":{"trusted":true},"cell_type":"code","source":"def image_names(partition='train'):\n    return os.listdir(f'/kaggle/input/pku-autonomous-driving/{partition}_images/')\n\ndef load_img(img_name, partition='train'):\n    image = Image.open(f'/kaggle/input/pku-autonomous-driving/{partition}_images/{img_name}')\n    image_mask = Image.open(f'/kaggle/input/pku-autonomous-driving/{partition}_masks/{img_name}')\n    return image, image_mask\n\ndef get_pred_string(img_name):\n    name = img_name.split('.')[0]\n    coords_str = TRAIN_DF[TRAIN_DF['ImageId'] == name]['PredictionString'].iloc[0]\n    return coords_str\n    \ndef add_mask(img, mask):\n    mask_array = np.array(mask)\n    masked = np.array(img)\n    masked[mask_array > 0] = 255\n    return Image.fromarray(masked)\n\ndef parse_pred_string(string):\n    \"\"\" Converts an string to a list of arrays with the positions of the cars. \"\"\"\n    cars = []\n    split_string = string.split()\n    for i in range(0, len(split_string), 7):\n        arr_i = split_string[i:i+7]\n        arr_i = [float(i) for i in arr_i]\n        arr_i[0] = int(arr_i[0])\n        cars.append(arr_i)\n    return cars\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We need a function to find the projections of the real world point (X, Y, Z) on the image (x, y). We can easily find them by multiplying the real world coordinates by the intrinsic matrix of the camera (which are provided inside the `camera_intrinsic.txt` file). \n\nThe f values are related to the camera sensor size (it's almost squared) whileas the c values correspond to the offset of the pixels from the origin point of the sensor (which in this case is located on the top left corner)"},{"metadata":{"trusted":true},"cell_type":"code","source":"CAMERA_fx = 2304.5479\nCAMERA_fy = 2305.8757\nCAMERA_cx = 1686.2379\nCAMERA_cy = 1354.9849\n\nCAMERA_MATRIX = np.array([\n    [CAMERA_fx,   0,        0],\n    [0,        CAMERA_fy,   0],\n    [0,           0,        1],\n])\n\ndef to_cam_xy(world_coords):\n    \"\"\" Converts world coordinates (X, Y, Z) to the projection on the images (x, y)\"\"\"\n    p = np.array(world_coords)\n    im_point = np.dot(p, CAMERA_MATRIX)\n    im_point[:,0] /= p[:,2]\n    im_point[:,1] /= p[:,2]\n    \n    im_point[:,0] += CAMERA_cx\n    im_point[:,1] += CAMERA_cy\n    \n    return im_point","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The next thing is to be able of plotting things arround the center of the car (i.e. bounding boxes, velocity vectors, etc.). For that, we will use the pitch, yaw and roll from the data. Using these values we will generate a rotation matrix which will rotate all the points we want.\n\nUsing this idea, we will use the `get_point_arround` function to add some points arround a `world_point` (which corresponds to the center of a car)  and then, rotate them taking as origin that `world_point`. Thus, if we have $p = <0,5,2>$, we can get a \"direction vector\" by adding some units to the $Z$ component and rotating it: $\\vec{v} = <0,5,(2+4)> * R$. \n\nOnce we have all the points we want, we simply re-project them using the `to_cam_xy` function."},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_rot_matrix(euler_rot):\n    yaw, pitch, roll = euler_rot\n    \n    yaw, pitch, roll = -yaw, -pitch, -roll\n\n    # The data reference edges seem to be rotated. This matrices work.\n    # I got the idea of flipping thanks to: \n    # https://www.kaggle.com/zstusnoopy/visualize-the-location-and-3d-bounding-box-of-car#kln-87\n    \n    rot_x = np.array([\n                        [1,     0,              0         ],\n                        [0, math.cos(yaw), -math.sin(yaw) ],\n                        [0, math.sin(yaw), math.cos(yaw)  ]\n                    ])\n         \n    rot_y = np.array([\n                        [math.cos(pitch),  0,      math.sin(pitch) ],\n                        [0,                1,      0               ],\n                        [-math.sin(pitch), 0,      math.cos(pitch) ]\n                    ])\n                 \n    rot_z = np.array([\n                        [math.cos(roll), -math.sin(roll), 0],\n                        [math.sin(roll),  math.cos(roll), 0],\n                        [0,               0,              1]\n                    ])\n                     \n                     \n    rotation_matrix = np.dot(rot_x, np.dot(rot_y, rot_z))\n \n    return rotation_matrix\n\ndef get_point_arround(world_point, rotation_angles, offsets=[[0,0,2]]):\n    \"\"\"Adds points arround the center (world point) and rotates them to match the \n    car rotation (taking as origin world point).\n    This can be used to calculate several points arround a vehicle (draw 3D bounding boxes, etc). \n    \n    Params:\n    world_point: numpy array [3] (x,y,z) of the car in the world.\n    rotation_angles: numpy array [3]. (yaw, pitch, roll) in radians.\n    offsets: List[List[3]] Points arround the world point (by default it is only one point 2 units ahead of the vehicle center).\n    \n    Returns:\n        Numpy array with all the point(s) rotated acordingly to the center point.\n    \"\"\"\n    rot_from_origin = np.eye(4)\n    origin = world_point\n    rot_from_origin[:3, 3] = origin\n    rot_from_origin[:3, :3] = get_rot_matrix(rotation_angles)\n    rot_from_origin = rot_from_origin[:3, :]\n    \n    points = np.ones((len(offsets), 4))\n    points[:,:3] = np.array(offsets)\n    points = points.T\n            \n    point = np.dot(rot_from_origin, points).T\n    return point\n    \n\n    \ndef plot_car_directions(image, coords):\n    \"\"\" Plots a point on each car and a green arrow pointing towards its direction (yaw pitch roll) \n    \n    Parameters:\n    image: PIL Image\n    coords: Coordinate array of the cars (this is the parsed string from the dataset).\n    \"\"\"\n    im = np.array(image)\n    \n    world_coords = [x[-3:] for x in coords]\n    courses = [x[1:4] for x in coords]\n\n    transformed = to_cam_xy(world_coords)\n\n    points_directions = [get_point_arround(center_point, rotaton) for center_point, rotaton in zip(world_coords, courses)]\n    transformed_dests = np.array([to_cam_xy(points) for points in points_directions])\n    \n    # Car position (in 2D)\n    x0 = transformed[:,0]\n    y0 = transformed[:,1]\n\n    # Movement vector (yaw, pitch, roll) in 2D\n    x1 = transformed_dests[:,:,0].flatten()\n    y1 = transformed_dests[:,:,1].flatten()\n    \n    for i in range(len(world_coords)):\n        im = cv2.arrowedLine(im, (int(x0[i]),int(y0[i])), (int(x1[i]),int(y1[i])), (0,255,0), 3, tipLength=0.06)\n        im = cv2.circle(im, (int(x0[i]),int(y0[i])), 10, (255,0,0), -1)\n    return Image.fromarray(im)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Loading data**"},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_IMAGES = image_names('train')\nTRAIN_DF = pd.read_csv('/kaggle/input/pku-autonomous-driving/train.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's test the above functions by loading several train images and plotting the car possitions and their direction vectors."},{"metadata":{"trusted":true},"cell_type":"code","source":"num_images = 4\nwidth = 2\n\nfig, ax = plt.subplots(num_images//width, width, figsize=(39,30))\n\nfor i in range(num_images):\n    test_image_name = TRAIN_IMAGES[400 + i]\n    image, image_mask = load_img(test_image_name)\n\n    coords_str = get_pred_string(test_image_name)\n    coords_str = parse_pred_string(coords_str)\n\n    world_coords = [x[-3:] for x in coords_str]\n    courses = [x[1:4] for x in coords_str]\n\n    transformed = to_cam_xy(world_coords)\n    image = plot_car_directions(image, coords_str)\n    ax[i//width, i%width].imshow(image)\n    #ax[i//width, i%width].set_aspect('equal')\n    ax[i//width, i%width].axis('off')\nplt.subplots_adjust(wspace=0, hspace=0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I've noticed some strange things. For example, in some photos, the direction vectors seem wrong for some vehicles. Maybe this dataset has been somehow auto-generated. We'll worry about this later."}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}