import matplotlib.pyplot as plt
import numpy as np
import cv2
import pandas as pd
import altair as alt


def main():
    # read label
    label_path = '' # replace with the label file generated by DeepLabCut network
    label = list()
    with open(label_path, 'r') as fin:
        idx = 0
        for line in fin:
            if idx > 2:
                line = line.strip().split(',')
                label.append(line[1:])
            idx += 1

    # read frames
    video_path = '' # the original video to analyze
    colors = [(255, 0, 0), (204,0,102), (255,51,153),(255,0,255), (255,255,102), (102,103,255), (102,51,153), (153,204,255), (0,255,255)]
    cap = cv2.VideoCapture(video_path)
    idx = 0
    idx_selected = 5700
    #idx_selected = 9420
    idx_selected = 4260
    prev_loc = dict()
    lines_to_draw = list()
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret == True:
            if idx_selected <= idx and idx <= idx_selected+30:
                if idx == idx_selected:
                    img = frame[:,:,:]
                if (idx-idx_selected)%5 == 0:
                    for i in range(9):
                        p = float(label[idx][3*i+2])
                        if p > 0.6:
                            location = (int(float(label[idx][3*i])), int(float(label[idx][3*i+1])))
                            color = colors[i]
                            img = cv2.circle(img, location, 10, color, thickness=-1)
                            
                            if i in prev_loc:
                                #img = cv2.line(img, prev_loc[i], location, (155,155,155), 3)
                                lines_to_draw.append((prev_loc[i], location))
                            prev_loc[i] = location

            idx += 1
        else:
            break
    for src, dst in lines_to_draw:
        img = cv2.line(img, src, dst, (45,45,45), 3)


if __name__ == 'main':
    main()

import matplotlib.pyplot as plt
import numpy as np
import cv2
import pandas as pd
import altair as alt


def main():
    # read label
    label_path = '' # replace with the label file generated by DeepLabCut network
    label = list()
    with open(label_path, 'r') as fin:
        idx = 0
        for line in fin:
            if idx > 2:
                line = line.strip().split(',')
                label.append(line[1:])
            idx += 1

    # read frames
    video_path = '' # the original video to analyze
    colors = [(255, 0, 0), (204,0,102), (255,51,153),(255,0,255), (255,255,102), (102,103,255), (102,51,153), (153,204,255), (0,255,255)]
    cap = cv2.VideoCapture(video_path)
    idx = 0
    idx_selected = 5700
    #idx_selected = 9420
    idx_selected = 4260
    prev_loc = dict()
    lines_to_draw = list()
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret == True:
            if idx_selected <= idx and idx <= idx_selected+30:
                if idx == idx_selected:
                    img = frame[:,:,:]
                if (idx-idx_selected)%5 == 0:
                    for i in range(9):
                        p = float(label[idx][3*i+2])
                        if p > 0.6:
                            location = (int(float(label[idx][3*i])), int(float(label[idx][3*i+1])))
                            color = colors[i]
                            img = cv2.circle(img, location, 10, color, thickness=-1)
                            
                            if i in prev_loc:
                                #img = cv2.line(img, prev_loc[i], location, (155,155,155), 3)
                                lines_to_draw.append((prev_loc[i], location))
                            prev_loc[i] = location

            idx += 1
        else:
            break
    for src, dst in lines_to_draw:
        img = cv2.line(img, src, dst, (45,45,45), 3)


if __name__ == 'main':
    main()

# read label # replace with the label file generated by DeepLabCut network # read frames # the original video to analyze #idx_selected = 9420 #img = cv2.line(img, prev_loc[i], location, (155,155,155), 3)