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

# # Compute (DLC) Pose Estimation accuracy
# First, we'll import the hand labeled ground truth files.
#
# This notebook is designed to work with the standard folder structure found in [DeepLabCut](https://github.com/DeepLabCut/DeepLabCut) projects. We are going to be using our own custom functions for evaluation here, so we're going to use DLC-Live so we can point to individual images rather than entire datasets with a subsection of labels.

# +
from dlclive import DLCLive, Processor
import cv2
from tqdm.notebook import tqdm, trange

"""
To install tqdm (to display progress bars in jupyter) run the following commands in your environment:

1. conda install -c conda-forge tqdm   # conda
2. pip install ipywidgets
3. jupyter nbextension enable --py widgetsnbextension

"""

dlc_proc = Processor()
dlc_trained_model_path = "C:/Users/Legos/Documents/PhD/Blender/OmniTrax/OmniTrax_WIP/DLC/DLC_multi_stick_v1_resnet_50_iteration-0_shuffle-1/"
dlc_live = DLCLive(dlc_trained_model_path, processor=dlc_proc)
# -

# now let's load all the data we want to evaluate!
#
# This script needs to point at the location of the "labeled-data" folder of your DLC project, expecting the hierachy:
#
# ***labeled-data/VIDEOS/CollectedData_USER.h5***

# +
import os
from pathlib import Path
import pandas as pd
labeled_data_path = Path("C:/Users/Legos/Documents/PhD/Blender/OmniTrax/OmniTrax_WIP/benchmark/STICK_DLC/")

# grab all h5 files
first_h5_found = False
for subdir, dirs, files in os.walk(labeled_data_path):
    for file in files:
        if file[-1] == "5":
            if not first_h5_found:
                df_labeled = pd.read_hdf(os.path.join(subdir, file))
                first_h5_found = True
            else:
                more_labeled = pd.read_hdf(os.path.join(subdir, file))
                df_labeled = df_labeled.append(more_labeled)

df_labeled

# +
# let's check the labels of the imported annotated data
import numpy as np
point_size = 5

for item in df_labeled.iterrows():
    img_name = item[0]
    img_path = os.path.join(labeled_data_path,img_name)
    img = cv2.imread(str(img_path))

    for coord in range(int((len(item[1])) / 2)):
        if not np.isnan(item[1][coord*2]) and not np.isnan(item[1][coord*2 + 1]):
            img = cv2.circle(img, (int(item[1][coord*2]), int(item[1][coord*2 + 1])),
                                       point_size,
                                       (int(255 * coord / 49), int(255 - 255 * coord / 49), 200), -1)
                   
    cv2.imshow("DLC Labeled IMG",img)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    
cv2.destroyAllWindows()
# -

# **Initialise** the network with the desired input shape. Afterwards the network can be used for inference on the fly

# +
desired_dims = (1024,768)

dlc_test_image_path = img_path

dlc_test_image = cv2.imread(dlc_test_image_path)
dlc_test_image_resized = cv2.resize(dlc_test_image,desired_dims)

dlc_live.init_inference(dlc_test_image_resized)

scale_factor = (desired_dims[0] / dlc_test_image.shape[1] + 
                desired_dims[1] / dlc_test_image.shape[0]) /2

print("Running inference at scale factor : ", scale_factor)
# -

# Cool, now let's check how well our trained network performs.
#
# If desired, add a lookup-table to only evaluate performance on specific landmarks. This can be useful when there is a discrepancy between which landmarks have been labeled in the real vs synthetic datasets.

# +
thresh = 0.5
point_size = 2

all_poses = [] # estimated landmark positions
all_scaled_poses = [] # rescaled to original resolution of labeled data
show_outputs = True

# optional landmark lookup table. denote skips with -1
"""
ID_lookup = [61,55,49,0,-1,-1,-1,3,5,-1,28,29,30,31,32,33,
             35,36,37,38,39,40,42,43,44,45,46,47,
             7,8,9,10,11,12,14,15,16,17,18,19,21,22,23,24,25,26]
"""             
ID_lookup = [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
             29,30,31,32,33,34,
             36,37,38,39,40,41,
             43,44,45,46,47,48,
             8,9,10,11,12,13,
             15,16,17,18,19,20,
             22,23,24,25,26,27]

pose_error = np.ones((len(df_labeled.index),len(ID_lookup))) * -1

print("Running inference...")

for i, item in tqdm(enumerate(df_labeled.iterrows()),total=len(df_labeled.index)):
    img_name = item[0]
    img_path = os.path.join(labeled_data_path,img_name)
    
    dlc_image = cv2.imread(img_path)
    dlc_image_resized = cv2.resize(dlc_image,desired_dims)

    pose = dlc_live.get_pose(dlc_image_resized)
    
    scaled_pose = pose.copy()
    for p, point in enumerate(pose):

        if show_outputs:
            if point[2] >= thresh:
                dlc_test_image_resized = cv2.circle(dlc_image_resized, (int(point[0]), int(point[1])),
                                           point_size,
                                           (int(255 * p / 49), int(255 - 255 * p / 49), 200), -1)
        
        # the shape of the image is returned in y , x, color
        p_x_scaled = (point[0] / desired_dims[0]) * dlc_image.shape[1]
        p_y_scaled = (point[1] / desired_dims[1]) * dlc_image.shape[0]

        scaled_pose[p][0:2] = p_x_scaled, p_y_scaled
        
        if p in ID_lookup and ID_lookup.index(p) != -1:
            if not pd.isnull(df_labeled.iloc[i,ID_lookup.index(p)*2]):
                if point[2] >= thresh:
                    #rint(scaled_pose[p][0:2],df_labeled.iloc[i,ID_lookup.index(p)*2], df_labeled.iloc[i,ID_lookup.index(p)*2 + 1])
                    pose_error[i,ID_lookup.index(p)] = np.linalg.norm(scaled_pose[p][0:2] - np.array([df_labeled.iloc[i,ID_lookup.index(p)*2], df_labeled.iloc[i,ID_lookup.index(p)*2 + 1]]))
        
    if show_outputs:
        cv2.imshow("DLC Predicted IMG",dlc_image_resized)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            show_outputs = False
            cv2.destroyAllWindows()
    
    all_poses.append(pose)
    all_scaled_poses.append(scaled_pose)
    
cv2.destroyAllWindows()

# ignore unassigned values when computing pose error
pose_error = np.where(pose_error == -1, np.nan, pose_error)
mean_pose_error = np.nanmean(pose_error)

print("\nMean Pose Error [px] :", round(mean_pose_error,3), " (at ground truth scale)")
print("Mean Pose Error [px] :", round(mean_pose_error * scale_factor,3), " (at inference scale)")

# -

per_part_MPE = np.nanmean(pose_error * scale_factor,axis=0)
for k, key in enumerate(df_labeled.keys()):
    if k % 2 == 0:
        if not ID_lookup[int(k / 2)] == -1:
            print(int(k / 2), key[1], abs(round(per_part_MPE[int(k / 2)],2)))



StarcoderdataJupyterScriptsDedupFiltered