Source code for simba.ui.pop_ups.directing_other_animals_plot_pop_up

__author__ = "Simon Nilsson; sronilsson@gmail.com"

import os
import threading
from tkinter import *
from typing import Union

from simba.mixins.config_reader import ConfigReader
from simba.mixins.pop_up_mixin import PopUpMixin
from simba.plotting.directing_animals_visualizer import \
    DirectingOtherAnimalsVisualizer
from simba.plotting.directing_animals_visualizer_mp import \
    DirectingOtherAnimalsVisualizerMultiprocess
from simba.ui.tkinter_functions import (CreateLabelFrameWithIcon, DropDownMenu,
                                        SimbaButton, SimBADropDown)
from simba.utils.enums import Formats, Keys, Links
from simba.utils.errors import AnimalNumberError, CountError
from simba.utils.lookups import find_closest_string, get_color_dict
from simba.utils.read_write import (find_all_videos_in_directory,
                                    find_core_cnt, str_2_bool)

DIRECTION_THICKNESS = "direction_thickness"
DIRECTIONALITY_COLOR = "directionality_color"
CIRCLE_SIZE = "circle_size"
HIGHLIGHT_ENDPOINTS = "highlight_endpoints"
SHOW_POSE = "show_pose"
ANIMAL_NAMES = "animal_names"
AUTO = 'AUTO'

NOSE, EAR_LEFT, EAR_RIGHT = Keys.NOSE.value, Keys.EAR_LEFT.value, Keys.EAR_RIGHT.value


[docs]class DirectingOtherAnimalsVisualizerPopUp(PopUpMixin, ConfigReader): def __init__(self, config_path: Union[str, os.PathLike]): ConfigReader.__init__(self, config_path=config_path, read_video_info=False) if self.animal_cnt == 1: raise AnimalNumberError(msg="Cannot visualize directionality between animals in a 1 animal project.", source=self.__class__.__name__,) PopUpMixin.__init__(self, title="CREATE ANIMAL DIRECTION VIDEOS", icon='eye') bp_names = list(set([x[:-2] for x in self.body_parts_lst])) nose_guess = find_closest_string(target=NOSE, string_list=bp_names)[0] ear_left_guess = find_closest_string(target=EAR_LEFT, string_list=bp_names)[0] ear_right_guess = find_closest_string(target=EAR_RIGHT, string_list=bp_names)[0] self.bp_frm = CreateLabelFrameWithIcon(parent=self.main_frm, header="SELECT BODY-PARTS", icon_name='pose', icon_link=Links.DIRECTING_ANIMALS_PLOTS.value) self.ear_left_dropdown = SimBADropDown(parent=self.bp_frm, dropdown_options=bp_names, label_width=35, dropdown_width=25, value=ear_left_guess, label='LEFT EAR body-part name:', img='ear_small', tooltip_key='DIRECTING_ANIMALS_LEFT_EAR') self.ear_right_dropdown = SimBADropDown(parent=self.bp_frm, dropdown_options=bp_names, label_width=35, dropdown_width=25, value=ear_right_guess, label='RIGHT EAR body-part name:', img='ear_right', tooltip_key='DIRECTING_ANIMALS_RIGHT_EAR') self.nose_dropdown = SimBADropDown(parent=self.bp_frm, dropdown_options=bp_names, label_width=35, dropdown_width=25, value=nose_guess, label='NOSE body-part name:', img='nose', tooltip_key='DIRECTING_ANIMALS_NOSE') self.color_dict = get_color_dict() self.color_lst = list(self.color_dict.keys()) self.color_lst.insert(0, "random") self.size_lst = list(range(1, 11)) self.size_lst.insert(0, AUTO) opacity_options = [round(x * 0.1, 1) for x in range(1, 11)] core_count = find_core_cnt()[0] self.files_found_dict = find_all_videos_in_directory(directory=self.video_dir, as_dict=True) self.style_settings_frm = CreateLabelFrameWithIcon(parent=self.main_frm, header="STYLE SETTINGS", icon_name=Keys.DOCUMENTATION.value, icon_link=Links.DIRECTING_ANIMALS_PLOTS.value,) self.show_pose_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=['TRUE', 'FALSE'], label="SHOW POSE TRACKING", label_width=35, dropdown_width=20, value="TRUE", img='pose', tooltip_key='DIRECTING_ANIMALS_SHOW_POSE') self.highlight_direction_endpoints_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=['TRUE', 'FALSE'], label="HIGHLIGHT DIRECTION END-POINTS:", label_width=35, dropdown_width=20, value="FALSE", img='finish', tooltip_key='DIRECTING_ANIMALS_HIGHLIGHT_ENDPOINTS') self.show_animal_names_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=['TRUE', 'FALSE'], label="SHOW ANIMAL NAMES:", label_width=35, dropdown_width=20, value="FALSE", img='id_card_2', tooltip_key='DIRECTING_ANIMALS_SHOW_NAMES') self.direction_clr_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=self.color_lst, label='DIRECTION COLOR:', label_width=35, dropdown_width=20, value="random", img='color_wheel', tooltip_key='DIRECTING_ANIMALS_DIRECTION_COLOR') self.pose_size_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=self.size_lst, label="POSE CIRCLE SIZE:", label_width=35, dropdown_width=20, value=AUTO, img='circle_small', tooltip_key='DIRECTING_ANIMALS_POSE_SIZE') self.line_thickness = SimBADropDown(parent=self.style_settings_frm, dropdown_options=self.size_lst, label="LINE THICKNESS:", label_width=35, dropdown_width=20, value=AUTO, img='line', tooltip_key='DIRECTING_ANIMALS_LINE_THICKNESS') self.line_opacity_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=opacity_options, label="LINE OPACITY:", label_width=35, dropdown_width=20, value=1.0, img='opacity', tooltip_key='DIRECTING_ANIMALS_LINE_OPACITY') self.core_count_dropdown = SimBADropDown(parent=self.style_settings_frm, dropdown_options=list(range(2, core_count+1)), label="CPU core count:", label_width=35, dropdown_width=20, value=int(core_count/3), img='cpu_small', tooltip_key='DIRECTING_ANIMALS_CPU_CORES') self.run_frm = CreateLabelFrameWithIcon(parent=self.main_frm, header="RUN", icon_name='rocket', icon_link=Links.DIRECTING_ANIMALS_PLOTS.value) self.run_single_video_frm = LabelFrame( self.run_frm, text="SINGLE VIDEO", font=Formats.FONT_HEADER.value, pady=5, padx=5, fg="black",) self.run_single_video_btn = SimbaButton(parent=self.run_single_video_frm, txt="Create single video", font=Formats.FONT_REGULAR.value, cmd=self.__create_directionality_plots, cmd_kwargs={'multiple_videos': False}, img='video_2') self.single_video_dropdown = SimBADropDown(parent=self.run_single_video_frm, dropdown_options=list(self.files_found_dict.keys()), label="VIDEO:", label_width=20, dropdown_width=20, value=list(self.files_found_dict.keys())[0], img='video_2', tooltip_key='DIRECTING_ANIMALS_SINGLE_VIDEO') self.run_multiple_videos = LabelFrame(self.run_frm,text="MULTIPLE VIDEO",font=Formats.FONT_HEADER.value,pady=5,padx=5,fg="black",) self.run_multiple_video_btn = SimbaButton(parent=self.run_multiple_videos, txt=f"Create multiple videos ({len(list(self.files_found_dict.keys()))} video(s) found)", font=Formats.FONT_REGULAR.value, cmd=self.__create_directionality_plots, cmd_kwargs={'multiple_videos': True}, img='stack') self.bp_frm.grid(row=0, column=0, sticky=NW) self.ear_left_dropdown.grid(row=0, column=0, sticky=NW) self.ear_right_dropdown.grid(row=1, column=0, sticky=NW) self.nose_dropdown.grid(row=2, column=0, sticky=NW) self.style_settings_frm.grid(row=1, column=0, sticky=NW) self.show_pose_dropdown.grid(row=0, column=0, sticky=NW) self.highlight_direction_endpoints_dropdown.grid(row=1, column=0, sticky=NW) self.show_animal_names_dropdown.grid(row=2, column=0, sticky=NW) self.direction_clr_dropdown.grid(row=3, column=0, sticky=NW) self.pose_size_dropdown.grid(row=4, column=0, sticky=NW) self.line_thickness.grid(row=5, column=0, sticky=NW) self.line_opacity_dropdown.grid(row=6, column=0, sticky=NW) self.core_count_dropdown.grid(row=7, column=0, sticky=NW) self.run_frm.grid(row=2, column=0, sticky=NW) self.run_single_video_frm.grid(row=0, column=0, sticky=NW) self.run_single_video_btn.grid(row=0, column=0, sticky=NW, padx=(0, 15)) self.single_video_dropdown.grid(row=0, column=1, sticky=NW) self.run_multiple_videos.grid(row=1, column=0, sticky=NW) self.run_multiple_video_btn.grid(row=0, column=0, sticky=NW) self.main_frm.mainloop() def __create_directionality_plots(self, multiple_videos: bool): if multiple_videos: video_paths = list(self.files_found_dict.values()) else: video_paths = [self.files_found_dict[self.single_video_dropdown.getChoices()]] show_pose = str_2_bool(self.show_pose_dropdown.get_value()) circle_size = None if self.pose_size_dropdown.get_value() == AUTO else int(self.pose_size_dropdown.get_value()) thickness = None if self.line_thickness.get_value() == AUTO else int(self.line_thickness.get_value()) show_names = str_2_bool(self.show_animal_names_dropdown.get_value()) highlight = str_2_bool(self.highlight_direction_endpoints_dropdown.get_value()) core_cnt = int(self.core_count_dropdown.get_value()) direction_clr = self.direction_clr_dropdown.getChoices() if direction_clr != "random": direction_clr = self.color_dict[direction_clr] opacity = float(self.line_opacity_dropdown.get_value()) nose = self.nose_dropdown.get_value() left_ear = self.ear_left_dropdown.get_value() right_ear = self.ear_right_dropdown.get_value() if len(list(set(list([nose, left_ear, right_ear])))) != 3: raise CountError(msg=f'The three chosen body-parts have to be unique: Got {nose, left_ear, right_ear}', source=self.__class__.__name__) style_attr = {SHOW_POSE: show_pose, CIRCLE_SIZE: circle_size, DIRECTIONALITY_COLOR: direction_clr, DIRECTION_THICKNESS: thickness, HIGHLIGHT_ENDPOINTS: highlight, ANIMAL_NAMES: show_names} for video_path in video_paths: if core_cnt == 1: visualizer = DirectingOtherAnimalsVisualizer(config_path=self.config_path, video_path=video_path, style_attr=style_attr, left_ear_name=left_ear, right_ear_name=right_ear, nose_name=nose) else: visualizer = DirectingOtherAnimalsVisualizerMultiprocess(config_path=self.config_path, video_path=video_path, style_attr=style_attr, core_cnt=core_cnt, left_ear_name=left_ear, line_opacity=opacity, right_ear_name=right_ear, nose_name=nose) threading.Thread(target=visualizer.run()).start()
#_ = DirectingOtherAnimalsVisualizerPopUp(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini') #_ = DirectingOtherAnimalsVisualizerPopUp(config_path='/Users/simon/Desktop/envs/troubleshooting/Two_animals_16bps/project_folder/project_config.ini') #_ = DirectingOtherAnimalsVisualizerPopUp(config_path=r"D:\troubleshooting\two_animals_sleap\project_folder\project_config.ini")