Source code for autopilot.transform.image

import os
import sys
import typing
from datetime import datetime

import numpy as np

from autopilot import prefs
from autopilot.transform.transforms import Transform

[docs]class Image(Transform): """ Metaclass for transformations of images """ def __init__(self, shape=None, *args, **kwargs): super(Image, self).__init__(*args, **kwargs) self._shape = shape # type: typing.Tuple[int, int] @property def format_in(self) -> dict: return { 'type': np.ndarray, 'shape': self.shape } @format_in.setter def format_in(self, format_in: dict): if 'shape' in format_in.keys(): self.shape = format_in['shape'] @property def format_out(self) -> dict: raise NotImplementedError('Every subclass of Transform must define format_out!') @format_out.setter def format_out(self, format_out: dict): raise NotImplementedError('Every subclass of Transform must define format_out!') @property def shape(self) -> typing.Tuple[int, int]: return self._shape @shape.setter def shape(self, shape): if not isinstance(shape, tuple) and not len(shape) == 2: raise ValueError('shape must be a tuple of two integers') self._shape = (round(shape[0]), round(shape[1]))
[docs]class DLC(Image): """ Do pose estimation with ``DeepLabCut-Live``!!!!! Specify a ``model_dir`` (relative to ``<BASEDIR>/dlc`` or absolute) or a model name from the DLC model zoo. All other args and kwargs are passed on to :class:`dlclive.DLCLive`, see its documentation for details: Attributes: model_type (str, 'local' or 'zoo'): whether a directory (local) or a modelzoo name (zoo) was passed live (:class:`dlclive.DLCLive`): the DLCLive object """ def __init__(self, model_dir: str = None, model_zoo: str = None, *args, **kwargs): """ Must give either ``model_dir`` or ``model_zoo`` Args: model_dir (str): directory of model, either absolute or relative to ``<BASEDIR>/dlc``. if ``None``, use ``model_zoo`` model_zoo (str): name of modelzoo model. if ``None``, use ``model_dir`` *args: passed to DLCLive and superclass **kwargs: passed to DLCLive and superclass """ super(DLC, self).__init__(*args, **kwargs) self._model = None self._model_dir = None self._deeplabcut = None self.model_type = None self.initialized = False if not model_dir and not model_zoo: raise ValueError('Either model_dir or model_zoo must be given!') elif model_dir and model_zoo: raise ValueError('Only model_dir OR model_zoo can be given!!') try: from dlclive import DLCLive except ImportError as e: print('dlclive not found! dlclive must be installed before using this transform') raise e self.import_dlc() self.deeplabcut = sys.modules['deeplabcut'] if model_dir: # figure out what model to use! self.model_type = 'local' self.model = model_dir self.model_dir = model_dir elif model_zoo: # check if valid self.model_type = "zoo" self.model = model_zoo # self.model_dir = os.path.join(self.dlc_dir, model_zoo) #makes sure model is exported, doesn't do anything if already exported self.export_model() = DLCLive(self.dlc_paths['export_dir'], *args, **kwargs)
[docs] def process(self, input: np.ndarray) -> np.ndarray: if not output = else: output = return output
@property def model(self) -> str: return self._model @model.setter def model(self, model: str): if self.model_type == 'zoo': if model not in self.list_modelzoo(): raise ValueError(f'model "{model}" not found in available models: {self.list_modelzoo()}') # check if we already have the model downloaded models = os.listdir(self.dlc_dir) models = [m for m in models if m.startswith(model)] if len(models) == 0: cfg_path = self.create_modelzoo(model) self.model_dir = os.path.dirname(cfg_path) elif len(models) == 1: self.model_dir = os.path.join(self.dlc_dir,models[0]) else: # more than one, pick the most recent one most_recent = datetime(1970,1,1) for i, test_model_dir in enumerate(models): pieces = test_model_dir.split('-') this_date = datetime(int(pieces[-3]), int(pieces[-2]), int(pieces[-1])) if this_date > most_recent: model_dir = test_model_dir most_recent = this_date self.model_dir = model_dir elif self.model_type == 'local': pass self._model = model @property def model_dir(self) -> str: return self._model_dir @model_dir.setter def model_dir(self, model_dir): if self.model_type == 'zoo': pass # if not os.path.exists(model_dir): # # # self.create_modelzoo() elif self.model_type == 'local': if not os.path.exists(model_dir): # if we were given a path that can be resolved, use it # otherwise, check if the model is in our DLC_DIR autopilot_model_dir = os.path.join(self.dlc_dir, model_dir) if os.path.exists(autopilot_model_dir): model_dir = autopilot_model_dir else: # see if a model *starting* with the model name exists autopilot_models = os.listdir(self.dlc_dir) autopilot_models = [a for a in autopilot_models if a.startswith(model_dir)] if len(autopilot_models)>0: model_dir = os.path.join(self.dlc_dir,autopilot_models[0]) if not os.path.exists(model_dir): raise ValueError(f'model_dir given, but model could not be found!\nmodel_dir: {model_dir}') self._model_dir = model_dir @property def dlc_paths(self) -> dict: """ paths used by dlc in manipulating/using models * config: <model_dir>/config.yaml * train_pose_cfg: <model_dir>/dlc-models/iteration-<n>/<name>/train/pose_cfg.yaml, * export_pose_cfg: <model_dir>/exported-models/<name>/pose_cfg.yaml * export_dir: <model_dir>/exported-models/<name> Returns: dict """ config = os.path.join(self.model_dir, 'config.yaml') # get pose_cfg in training folder train_pose_cfg = None # get latest iteration dlc_models = os.path.join(self.model_dir, 'dlc-models') iteration = sorted([it for it in os.listdir(dlc_models) if it.startswith('iteration')]) if len(iteration)>0: # within an iteration, maybe multiple training sessions iteration = os.path.join(dlc_models, iteration[-1]) training_session = sorted(os.listdir(iteration)) if len(training_session)>0: test_train_pose_cfg = os.path.join(iteration, training_session[-1], 'train', 'pose_cfg.yaml') if os.path.exists(test_train_pose_cfg): train_pose_cfg = test_train_pose_cfg # get exported pose_cfg.yaml exported_pose_cfg = None exported_dir = os.path.join(self.model_dir, 'exported-models') exported_model_dir = None if os.path.exists(exported_dir): # FIXME: assuming only one exported model for now exported_subdirs = sorted([subd for subd in os.listdir(exported_dir) if \ os.path.isdir(os.path.join(exported_dir,subd))]) test_exported_pose_cfg = os.path.join(exported_dir, exported_subdirs[-1], 'pose_cfg.yaml') if os.path.exists(test_exported_pose_cfg): exported_pose_cfg = test_exported_pose_cfg exported_model_dir = os.path.join(exported_dir, exported_subdirs[-1]) return { 'config':config, 'train_pose_cfg': train_pose_cfg, 'export_pose_cfg': exported_pose_cfg, 'export_dir': exported_model_dir } @property def dlc_dir(self) -> str: """ ``{prefs.get('BASE_DIR')}/dlc`` Returns: str """ if 'DLCDIR' in prefs._PREFS.keys(): dlc_dir = prefs.get('DLCDIR') else: dlc_dir = os.path.join(prefs.get('BASEDIR'), 'dlc') if not os.path.exists(dlc_dir): try: os.mkdir(dlc_dir) except OSError as e: raise OSError(f'No DLC dir found and one could not be created!\n{e}') prefs.set('DLC_DIR', dlc_dir) return dlc_dir
[docs] @classmethod def list_modelzoo(cls): """ List available modelzoo model names in local deeplabcut version Returns: list: names of available modelzoo models """ __deeplabcut = __import__('deeplabcut') return __deeplabcut.create_project.modelzoo.Modeloptions
# # @property # def deeplabcut(self): # if not self._deeplabcut: # try: # os.environ['DLClight'] = "True" # os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # self._deeplabcut = __import__('deeplabcut') # except ImportError as e: # print('deeplabcut not found! deeplabcut must be installed before using this transform!') # raise e # return self._deeplabcut
[docs] def import_dlc(cls): if 'deeplabcut' not in sys.modules.keys(): try: os.environ['DLClight'] = "True" os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import deeplabcut except ImportError as e: print('deeplabcut not found! deeplabcut must be installed before using this transform!') raise e
[docs] def create_modelzoo(self, model): self.model_type = "zoo" cfg_path, _ = self.deeplabcut.create_pretrained_project( model, 'autopilot', [], model=model, working_directory=self.dlc_dir, analyzevideo=False, createlabeledvideo=False, exportmodel=True) return cfg_path
[docs] def load_model(self): pass
[docs] def export_model(self): # pdb.set_trace() if not self.dlc_paths['export_pose_cfg']: try: self.deeplabcut.export_model(self.dlc_paths['config']) except FileExistsError: pass
@property def format_in(self) -> dict: return { 'type': np.ndarray } @property def format_out(self) -> dict: return { 'type': np.ndarray }