diff --git a/.gitignore b/.gitignore index 0c3e79a..e373334 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,9 @@ dlclive/check_install/dlc-live-tmp* **DS_Store* *vscode* +# Model zoo snapshots +dlclive/modelzoo/snapshots/* + **/__MACOSX/ # Include tests directory (negate the *test* pattern for tests/) diff --git a/dlclive/modelzoo/__init__.py b/dlclive/modelzoo/__init__.py new file mode 100644 index 0000000..1f4f018 --- /dev/null +++ b/dlclive/modelzoo/__init__.py @@ -0,0 +1,9 @@ +from dlclive.modelzoo.utils import ( + _MODELZOO_PATH, + list_available_models, + list_available_projects, + list_available_combinations, + load_super_animal_config, + download_super_animal_snapshot, +) +from dlclive.modelzoo.pytorch_model_zoo_export import export_modelzoo_model \ No newline at end of file diff --git a/dlclive/modelzoo/model_configs/dlcrnet.yaml b/dlclive/modelzoo/model_configs/dlcrnet.yaml new file mode 100644 index 0000000..570c1a0 --- /dev/null +++ b/dlclive/modelzoo/model_configs/dlcrnet.yaml @@ -0,0 +1,121 @@ + # Project definitions (do not edit) +Task: +scorer: +date: +multianimalproject: +identity: + + # Project path (change when moving around) +project_path: + + # Annotation data set configuration (and individual video cropping parameters) +video_sets: +bodyparts: + + # Fraction of video to start/stop when extracting frames for labeling/refinement +start: +stop: +numframes2pick: + + # Plotting configuration +skeleton: [] +skeleton_color: black +pcutoff: +dotsize: +alphavalue: +colormap: + + # Training,Evaluation and Analysis configuration +TrainingFraction: +iteration: +default_net_type: +default_augmenter: +snapshotindex: +batch_size: 1 + + # Cropping Parameters (for analysis and outlier frame detection) +cropping: + #if cropping is true for analysis, then set the values here: +x1: +x2: +y1: +y2: + + # Refinement configuration (parameters from annotation dataset configuration also relevant in this stage) +corner2move2: +move2corner: +alpha_r: 0.02 +apply_prob: 0.5 +clahe: true +claheratio: 0.1 +crop_sampling: hybrid +crop_size: +- 400 +- 400 +cropratio: 0.4 +dataset: +dataset_type: multi-animal-imgaug +decay_steps: 30000 +display_iters: 500 +edge: false +emboss: + alpha: + - 0.0 + - 1.0 + embossratio: 0.1 + strength: + - 0.5 + - 1.5 +global_scale: 0.8 +histeq: true +histeqratio: 0.1 +init_weights: +intermediate_supervision: false +intermediate_supervision_layer: 12 +location_refinement: true +locref_huber_loss: true +locref_loss_weight: 0.05 +locref_stdev: 7.2801 +lr_init: 0.0005 +max_input_size: 1500 +max_shift: 0.4 +mean_pixel: +- 123.68 +- 116.779 +- 103.939 +metadataset: +min_input_size: 64 +mirror: false +multi_stage: true +multi_step: +- - 0.0001 + - 7500 +- - 5.0e-05 + - 12000 +- - 1.0e-05 + - 200000 +net_type: resnet_50 +num_idchannel: 0 +num_joints: 27 +num_limbs: 351 +optimizer: adam +pafwidth: 20 +pairwise_huber_loss: false +pairwise_loss_weight: 0.1 +pairwise_predict: false +partaffinityfield_graph: [] +partaffinityfield_predict: false +pos_dist_thresh: 17 +pre_resize: [] +rotation: 25 +rotratio: 0.4 +save_iters: 10000 +scale_jitter_lo: 0.5 +scale_jitter_up: 1.25 +sharpen: false +sharpenratio: 0.3 +stride: 8.0 +weigh_only_present_joints: false +gradient_masking: true +weight_decay: 0.0001 +weigh_part_predictions: false diff --git a/dlclive/modelzoo/model_configs/fasterrcnn_mobilenet_v3_large_fpn.yaml b/dlclive/modelzoo/model_configs/fasterrcnn_mobilenet_v3_large_fpn.yaml new file mode 100644 index 0000000..6d4e11b --- /dev/null +++ b/dlclive/modelzoo/model_configs/fasterrcnn_mobilenet_v3_large_fpn.yaml @@ -0,0 +1,51 @@ +data: + colormode: RGB + inference: + normalize_images: true + train: + affine: + p: 0.5 + rotation: 30 + scaling: [ 1.0, 1.0 ] + translation: 40 + collate: + type: ResizeFromDataSizeCollate + min_scale: 0.4 + max_scale: 1.0 + min_short_side: 128 + max_short_side: 1152 + multiple_of: 32 + to_square: false + hflip: true + normalize_images: true +device: auto +model: + type: FasterRCNN + variant: fasterrcnn_mobilenet_v3_large_fpn + box_score_thresh: 0.6 + freeze_bn_stats: true + freeze_bn_weights: false +runner: + type: DetectorTrainingRunner + key_metric: "test.mAP@50:95" + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 1e-5 + scheduler: + type: LRListScheduler + params: + milestones: [ 90 ] + lr_list: [ [ 1e-6 ] ] + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 1 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 250 diff --git a/dlclive/modelzoo/model_configs/fasterrcnn_resnet50_fpn_v2.yaml b/dlclive/modelzoo/model_configs/fasterrcnn_resnet50_fpn_v2.yaml new file mode 100644 index 0000000..27d147e --- /dev/null +++ b/dlclive/modelzoo/model_configs/fasterrcnn_resnet50_fpn_v2.yaml @@ -0,0 +1,51 @@ +data: + colormode: RGB + inference: + normalize_images: true + train: + affine: + p: 0.5 + rotation: 30 + scaling: [ 1.0, 1.0 ] + translation: 40 + collate: + type: ResizeFromDataSizeCollate + min_scale: 0.4 + max_scale: 1.0 + min_short_side: 128 + max_short_side: 1152 + multiple_of: 32 + to_square: false + hflip: true + normalize_images: true +device: auto +model: + type: FasterRCNN + variant: fasterrcnn_resnet50_fpn_v2 + box_score_thresh: 0.6 + freeze_bn_stats: true + freeze_bn_weights: false +runner: + type: DetectorTrainingRunner + key_metric: "test.mAP@50:95" + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 1e-5 + scheduler: + type: LRListScheduler + params: + milestones: [ 90 ] + lr_list: [ [ 1e-6 ] ] + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 1 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 250 \ No newline at end of file diff --git a/dlclive/modelzoo/model_configs/hrnet_w32.yaml b/dlclive/modelzoo/model_configs/hrnet_w32.yaml new file mode 100644 index 0000000..011d772 --- /dev/null +++ b/dlclive/modelzoo/model_configs/hrnet_w32.yaml @@ -0,0 +1,81 @@ +data: + colormode: RGB + inference: + auto_padding: + pad_width_divisor: 32 + pad_height_divisor: 32 + normalize_images: true + train: + affine: + p: 0.5 + scaling: [1.0, 1.0] + rotation: 30 + translation: 0 + gaussian_noise: 12.75 + normalize_images: true + auto_padding: + pad_width_divisor: 32 + pad_height_divisor: 32 +device: auto +method: td +model: + backbone: + type: HRNet + model_name: hrnet_w32 + pretrained: false + freeze_bn_stats: True + freeze_bn_weights: False + interpolate_branches: false + increased_channel_count: false + backbone_output_channels: 32 + heads: + bodypart: + type: HeatmapHead + weight_init: "normal" + predictor: + type: HeatmapPredictor + apply_sigmoid: false + clip_scores: true + location_refinement: false + locref_std: 7.2801 + target_generator: + type: HeatmapGaussianGenerator + num_heatmaps: "num_bodyparts" + pos_dist_thresh: 17 + heatmap_mode: KEYPOINT + generate_locref: false + locref_std: 7.2801 + criterion: + heatmap: + type: WeightedMSECriterion + weight: 1.0 + heatmap_config: + channels: [32, "num_bodyparts"] + kernel_size: [1] + strides: [1] +net_type: hrnet_w32 +runner: + type: PoseTrainingRunner + key_metric: "test.mAP" + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 1e-5 + scheduler: + type: LRListScheduler + params: + lr_list: [ [ 1e-6 ], [ 1e-7 ] ] + milestones: [ 160, 190 ] + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 1 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 200 + seed: 42 diff --git a/dlclive/modelzoo/model_configs/resnet_50.yaml b/dlclive/modelzoo/model_configs/resnet_50.yaml new file mode 100644 index 0000000..994840c --- /dev/null +++ b/dlclive/modelzoo/model_configs/resnet_50.yaml @@ -0,0 +1,88 @@ +data: + colormode: RGB + inference: + normalize_images: true + train: + affine: + p: 0.5 + scaling: [1.0, 1.0] + rotation: 30 + translation: 0 + gaussian_noise: 12.75 + normalize_images: true +device: auto +method: td +model: + backbone: + type: ResNet + model_name: resnet50_gn + output_stride: 16 + freeze_bn_stats: false + freeze_bn_weights: false + backbone_output_channels: 2048 + heads: + bodypart: + type: HeatmapHead + weight_init: normal + predictor: + type: HeatmapPredictor + apply_sigmoid: false + clip_scores: true + location_refinement: true + locref_std: 7.2801 + target_generator: + type: HeatmapGaussianGenerator + num_heatmaps: "num_bodyparts" + pos_dist_thresh: 17 + heatmap_mode: KEYPOINT + generate_locref: true + locref_std: 7.2801 + criterion: + heatmap: + type: WeightedMSECriterion + weight: 1.0 + locref: + type: WeightedHuberCriterion + weight: 0.05 + heatmap_config: + channels: + - 2048 + - "num_bodyparts" + kernel_size: + - 3 + strides: + - 2 + locref_config: + channels: + - 2048 + - "num_bodyparts x 2" + kernel_size: + - 3 + strides: + - 2 +net_type: resnet_50 +runner: + type: PoseTrainingRunner + key_metric: "test.mAP" + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 1e-5 + scheduler: + type: LRListScheduler + params: + lr_list: [ [ 1e-6 ], [ 1e-7 ] ] + milestones: [ 160, 190 ] + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 1 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 100 + seed: 42 diff --git a/dlclive/modelzoo/model_configs/rtmpose_s.yaml b/dlclive/modelzoo/model_configs/rtmpose_s.yaml new file mode 100644 index 0000000..9c80b05 --- /dev/null +++ b/dlclive/modelzoo/model_configs/rtmpose_s.yaml @@ -0,0 +1,102 @@ +data: + colormode: RGB + inference: + normalize_images: true + train: + affine: + p: 0.5 + scaling: [1.0, 1.0] + rotation: 30 + translation: 0 + gaussian_noise: 12.75 + normalize_images: true +device: auto +method: td +model: + backbone: + type: CSPNeXt + model_name: cspnext_s + freeze_bn_stats: false + freeze_bn_weights: false + deepen_factor: 0.33 + widen_factor: 0.5 + backbone_output_channels: 512 + heads: + bodypart: + type: RTMCCHead + weight_init: RTMPose + target_generator: + type: SimCCGenerator + input_size: + - 256 + - 256 + smoothing_type: gaussian + sigma: + - 5.66 + - 5.66 + simcc_split_ratio: 2.0 + label_smooth_weight: 0.0 + normalize: false + criterion: + x: + type: KLDiscreteLoss + use_target_weight: true + beta: 10.0 + label_softmax: true + y: + type: KLDiscreteLoss + use_target_weight: true + beta: 10.0 + label_softmax: true + predictor: + type: SimCCPredictor + simcc_split_ratio: 2.0 + sigma: + - 5.66 + - 5.66 + decode_beta: 150.0 + input_size: + - 256 + - 256 + in_channels: 512 + out_channels: 39 + in_featuremap_size: + - 8 + - 8 + simcc_split_ratio: 2.0 + final_layer_kernel_size: 7 + gau_cfg: + hidden_dims: 256 + s: 128 + expansion_factor: 2 + dropout_rate: 0 + drop_path: 0.0 + act_fn: SiLU + use_rel_bias: false + pos_enc: false +net_type: rtmpose_s +runner: + type: PoseTrainingRunner + key_metric: "test.mAP" + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 1e-5 + scheduler: + type: LRListScheduler + params: + lr_list: [ [ 1e-6 ], [ 1e-7 ] ] + milestones: [ 160, 190 ] + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 1 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 100 + seed: 42 diff --git a/dlclive/modelzoo/model_configs/rtmpose_x.yaml b/dlclive/modelzoo/model_configs/rtmpose_x.yaml new file mode 100644 index 0000000..0d1fb8a --- /dev/null +++ b/dlclive/modelzoo/model_configs/rtmpose_x.yaml @@ -0,0 +1,163 @@ +data: + colormode: RGB + inference: + normalize_images: true + top_down_crop: + width: 288 + height: 384 + train: + affine: + p: 0.5 + rotation: 30 + scaling: + - 1.0 + - 1.0 + translation: 0 + collate: + covering: false + gaussian_noise: 12.75 + hist_eq: false + motion_blur: false + normalize_images: true + top_down_crop: + width: 288 + height: 384 +detector: null +device: auto +metadata: + project_path: null + pose_config_path: rtmpose_x_body7_pytorch_config.yaml + bodyparts: + - nose + - left_eye + - right_eye + - left_ear + - right_ear + - left_shoulder + - right_shoulder + - left_elbow + - right_elbow + - left_wrist + - right_wrist + - left_hip + - right_hip + - left_knee + - right_knee + - left_ankle + - right_ankle + unique_bodyparts: [] + individuals: + - idv0 + - idv1 + - idv2 + - idv3 + - idv4 + - idv5 + - idv6 + - idv7 + - idv8 + - idv9 + with_identity: false +method: td +model: + backbone: + type: CSPNeXt + model_name: cspnext_p5 + freeze_bn_stats: false + freeze_bn_weights: false + expand_ratio: 0.5 + deepen_factor: 1.33 + widen_factor: 1.25 + channel_attention: true + norm_layer: SyncBN + activation_fn: SiLU + backbone_output_channels: 1280 + heads: + bodypart: + type: RTMCCHead + weight_init: RTMPose + target_generator: + type: SimCCGenerator + input_size: + - 288 + - 384 + smoothing_type: gaussian + sigma: + - 6.0 + - 6.93 + simcc_split_ratio: 2.0 + label_smooth_weight: 0.0 + normalize: false + criterion: + x: + type: KLDiscreteLoss + use_target_weight: true + beta: 10.0 + label_softmax: true + y: + type: KLDiscreteLoss + use_target_weight: true + beta: 10.0 + label_softmax: true + predictor: + type: SimCCPredictor + simcc_split_ratio: 2.0 + sigma: + - 6.0 + - 6.93 + decode_beta: 150.0 + input_size: + - 288 + - 384 + in_channels: 1280 + out_channels: 17 + in_featuremap_size: + - 9 + - 12 + simcc_split_ratio: 2.0 + final_layer_kernel_size: 7 + gau_cfg: + hidden_dims: 256 + s: 128 + expansion_factor: 2 + dropout_rate: 0 + drop_path: 0.0 + act_fn: SiLU + use_rel_bias: false + pos_enc: false +net_type: rtmpose_x +runner: + type: PoseTrainingRunner + gpus: + key_metric: test.mAP + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 0.0005 + scheduler: + type: SequentialLR + params: + schedulers: + - type: ConstantLR + params: + factor: 0.001 + total_iters: 5 + - type: CosineAnnealingLR + params: + T_max: 250 + eta_min: 1e-05 + milestones: + - 100 + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 1 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 200 + seed: 42 diff --git a/dlclive/modelzoo/model_configs/ssdlite.yaml b/dlclive/modelzoo/model_configs/ssdlite.yaml new file mode 100644 index 0000000..04e694f --- /dev/null +++ b/dlclive/modelzoo/model_configs/ssdlite.yaml @@ -0,0 +1,50 @@ +data: + colormode: RGB + inference: + normalize_images: true + train: + affine: + p: 0.5 + rotation: 30 + scaling: [ 1.0, 1.0 ] + translation: 40 + collate: + type: ResizeFromDataSizeCollate + min_scale: 0.4 + max_scale: 1.0 + min_short_side: 128 + max_short_side: 1152 + multiple_of: 32 + to_square: false + hflip: true + normalize_images: true +device: auto +model: + type: SSDLite + box_score_thresh: 0.6 + freeze_bn_stats: true + freeze_bn_weights: false +runner: + type: DetectorTrainingRunner + key_metric: "test.mAP@50:95" + key_metric_asc: true + eval_interval: 10 + optimizer: + type: AdamW + params: + lr: 1e-5 + scheduler: + type: LRListScheduler + params: + milestones: [ 90 ] + lr_list: [ [ 1e-6 ] ] + snapshots: + max_snapshots: 5 + save_epochs: 25 + save_optimizer_state: false +train_settings: + batch_size: 8 + dataloader_workers: 0 + dataloader_pin_memory: false + display_iters: 500 + epochs: 250 \ No newline at end of file diff --git a/dlclive/modelzoo/project_configs/superanimal_bird.yaml b/dlclive/modelzoo/project_configs/superanimal_bird.yaml new file mode 100644 index 0000000..3144433 --- /dev/null +++ b/dlclive/modelzoo/project_configs/superanimal_bird.yaml @@ -0,0 +1,105 @@ +# Project definitions (do not edit) +Task: +scorer: +date: +multianimalproject: +identity: + + +# Project path (change when moving around) +project_path: + + +# Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow) +engine: pytorch + + +# Annotation data set configuration (and individual video cropping parameters) +video_sets: +bodyparts: +- back +- bill +- belly +- breast +- crown +- forehead +- left_eye +- left_leg +- left_wing_tip +- left_wrist +- nape +- right_eye +- right_leg +- right_wing_tip +- right_wrist +- tail_tip +- throat +- neck +- tail_left +- tail_right +- upper_spine +- upper_half_spine +- lower_half_spine +- right_foot +- left_foot +- left_half_chest +- right_half_chest +- chin +- left_tibia +- right_tibia +- lower_spine +- upper_half_neck +- lower_half_neck +- left_chest +- right_chest +- upper_neck +- left_wing_shoulder +- left_wing_elbow +- right_wing_shoulder +- right_wing_elbow +- upper_cere +- lower_cere + + +# Fraction of video to start/stop when extracting frames for labeling/refinement +start: +stop: +numframes2pick: + + +# Plotting configuration +skeleton: [] +skeleton_color: black +pcutoff: +dotsize: +alphavalue: +colormap: rainbow + + +# Training,Evaluation and Analysis configuration +TrainingFraction: +iteration: +default_net_type: +default_augmenter: +snapshotindex: +detector_snapshotindex: -1 +batch_size: 1 +detector_batch_size: 1 + + +# Cropping Parameters (for analysis and outlier frame detection) +cropping: +#if cropping is true for analysis, then set the values here: +x1: +x2: +y1: +y2: + + +# Refinement configuration (parameters from annotation dataset configuration also relevant in this stage) +corner2move2: +move2corner: + + +# Conversion tables to fine-tune SuperAnimal weights +SuperAnimalConversionTables: diff --git a/dlclive/modelzoo/project_configs/superanimal_humanbody.yaml b/dlclive/modelzoo/project_configs/superanimal_humanbody.yaml new file mode 100644 index 0000000..d1e665c --- /dev/null +++ b/dlclive/modelzoo/project_configs/superanimal_humanbody.yaml @@ -0,0 +1,90 @@ +# Project definitions (do not edit) +Task: +scorer: +date: +multianimalproject: +identity: + +# Project path (change when moving around) +project_path: + +# Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow) +engine: pytorch + +# Annotation data set configuration (and individual video cropping parameters) +video_sets: +bodyparts: +- nose +- left_eye +- right_eye +- left_ear +- right_ear +- left_shoulder +- right_shoulder +- left_elbow +- right_elbow +- left_wrist +- right_wrist +- left_hip +- right_hip +- left_knee +- right_knee +- left_ankle +- right_ankle + +# Fraction of video to start/stop when extracting frames for labeling/refinement +start: +stop: +numframes2pick: + +# Plotting configuration +skeleton: + - [16, 14] + - [14, 12] + - [17, 15] + - [15, 13] + - [12, 13] + - [6, 12] + - [7, 13] + - [6, 7] + - [6, 8] + - [7, 9] + - [8, 10] + - [9, 11] + - [2, 3] + - [1, 2] + - [1, 3] + - [2, 4] + - [3, 5] + - [4, 6] + - [5, 7] +skeleton_color: black +pcutoff: +dotsize: +alphavalue: +colormap: rainbow + +# Training,Evaluation and Analysis configuration +TrainingFraction: +iteration: +default_net_type: +default_augmenter: +snapshotindex: +detector_snapshotindex: -1 +batch_size: 1 +detector_batch_size: 1 + +# Cropping Parameters (for analysis and outlier frame detection) +cropping: +#if cropping is true for analysis, then set the values here: +x1: +x2: +y1: +y2: + +# Refinement configuration (parameters from annotation dataset configuration also relevant in this stage) +corner2move2: +move2corner: + +# Conversion tables to fine-tune SuperAnimal weights +SuperAnimalConversionTables: \ No newline at end of file diff --git a/dlclive/modelzoo/project_configs/superanimal_quadruped.yaml b/dlclive/modelzoo/project_configs/superanimal_quadruped.yaml new file mode 100644 index 0000000..e06f820 --- /dev/null +++ b/dlclive/modelzoo/project_configs/superanimal_quadruped.yaml @@ -0,0 +1,102 @@ +# Project definitions (do not edit) +Task: +scorer: +date: +multianimalproject: +identity: + + +# Project path (change when moving around) +project_path: + + +# Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow) +engine: pytorch + + +# Annotation data set configuration (and individual video cropping parameters) +video_sets: +bodyparts: +- nose +- upper_jaw +- lower_jaw +- mouth_end_right +- mouth_end_left +- right_eye +- right_earbase +- right_earend +- right_antler_base +- right_antler_end +- left_eye +- left_earbase +- left_earend +- left_antler_base +- left_antler_end +- neck_base +- neck_end +- throat_base +- throat_end +- back_base +- back_end +- back_middle +- tail_base +- tail_end +- front_left_thai +- front_left_knee +- front_left_paw +- front_right_thai +- front_right_knee +- front_right_paw +- back_left_paw +- back_left_thai +- back_right_thai +- back_left_knee +- back_right_knee +- back_right_paw +- belly_bottom +- body_middle_right +- body_middle_left + + +# Fraction of video to start/stop when extracting frames for labeling/refinement +start: +stop: +numframes2pick: + + +# Plotting configuration +skeleton: [] +skeleton_color: black +pcutoff: +dotsize: +alphavalue: +colormap: rainbow + + +# Training,Evaluation and Analysis configuration +TrainingFraction: +iteration: +default_net_type: +default_augmenter: +snapshotindex: +detector_snapshotindex: -1 +batch_size: 1 +detector_batch_size: 1 + + +# Cropping Parameters (for analysis and outlier frame detection) +cropping: +#if cropping is true for analysis, then set the values here: +x1: +x2: +y1: +y2: + + +# Refinement configuration (parameters from annotation dataset configuration also relevant in this stage) +corner2move2: +move2corner: + + +# Conversion tables to fine-tune SuperAnimal weights +SuperAnimalConversionTables: diff --git a/dlclive/modelzoo/project_configs/superanimal_topviewmouse.yaml b/dlclive/modelzoo/project_configs/superanimal_topviewmouse.yaml new file mode 100644 index 0000000..bd09628 --- /dev/null +++ b/dlclive/modelzoo/project_configs/superanimal_topviewmouse.yaml @@ -0,0 +1,90 @@ +# Project definitions (do not edit) +Task: +scorer: +date: +multianimalproject: +identity: + + +# Project path (change when moving around) +project_path: + + +# Default DeepLabCut engine to use for shuffle creation (either pytorch or tensorflow) +engine: pytorch + + +# Annotation data set configuration (and individual video cropping parameters) +video_sets: +bodyparts: +- nose +- left_ear +- right_ear +- left_ear_tip +- right_ear_tip +- left_eye +- right_eye +- neck +- mid_back +- mouse_center +- mid_backend +- mid_backend2 +- mid_backend3 +- tail_base +- tail1 +- tail2 +- tail3 +- tail4 +- tail5 +- left_shoulder +- left_midside +- left_hip +- right_shoulder +- right_midside +- right_hip +- tail_end +- head_midpoint + + +# Fraction of video to start/stop when extracting frames for labeling/refinement +start: +stop: +numframes2pick: + + +# Plotting configuration +skeleton: [] +skeleton_color: black +pcutoff: +dotsize: +alphavalue: +colormap: rainbow + + +# Training,Evaluation and Analysis configuration +TrainingFraction: +iteration: +default_net_type: +default_augmenter: +snapshotindex: +detector_snapshotindex: -1 +batch_size: 1 +detector_batch_size: 1 + + +# Cropping Parameters (for analysis and outlier frame detection) +cropping: +#if cropping is true for analysis, then set the values here: +x1: +x2: +y1: +y2: + + +# Refinement configuration (parameters from annotation dataset configuration also relevant in this stage) +corner2move2: +move2corner: + + +# Conversion tables to fine-tune SuperAnimal weights +SuperAnimalConversionTables: diff --git a/dlclive/modelzoo/pytorch_model_zoo_export.py b/dlclive/modelzoo/pytorch_model_zoo_export.py new file mode 100644 index 0000000..616857d --- /dev/null +++ b/dlclive/modelzoo/pytorch_model_zoo_export.py @@ -0,0 +1,54 @@ +import warnings +from pathlib import Path +from collections import OrderedDict + +import torch + +from dlclive.modelzoo.utils import load_super_animal_config, download_super_animal_snapshot + + +def export_modelzoo_model( + export_path: str | Path, + super_animal: str, + model_name: str, + detector_name: str | None = None, +) -> None: + """ + + """ + Path(export_path).parent.mkdir(parents=True, exist_ok=True) + if Path(export_path).exists(): + warnings.warn(f"Export path {export_path} already exists, skipping export", UserWarning) + return + + model_cfg = load_super_animal_config( + super_animal=super_animal, + model_name=model_name, + detector_name=detector_name, + ) + + def _load_model_weights(model_name: str, super_animal: str = super_animal) -> OrderedDict: + """Download the model weights from huggingface and load them in torch state dict""" + checkpoint: Path = download_super_animal_snapshot(dataset=super_animal, model_name=model_name) + return torch.load(checkpoint, map_location="cpu", weights_only=True)["model"] + + export_dict = { + "config": model_cfg, + "pose": _load_model_weights(model_name), + "detector": _load_model_weights(detector_name) if detector_name is not None else None, + } + torch.save(export_dict, export_path) + + +if __name__ == "__main__": + """Example usage""" + from utils import _MODELZOO_PATH + + model_name = "resnet_50" + super_animal = "superanimal_quadruped" + + export_modelzoo_model( + export_path=_MODELZOO_PATH / 'exported_models' / f'exported_{super_animal}_{model_name}.pt', + super_animal=super_animal, + model_name=model_name, + ) diff --git a/dlclive/modelzoo/resolve_config.py b/dlclive/modelzoo/resolve_config.py new file mode 100644 index 0000000..bea25f5 --- /dev/null +++ b/dlclive/modelzoo/resolve_config.py @@ -0,0 +1,136 @@ +""" +Helper function to deal with default values in the model configuration. +For instance, "num_bodyparts x 2" is replaced with the number of bodyparts multiplied by 2. +""" +# NOTE JR 2026-23-01: This is duplicate code, copied from the original DeepLabCut-Live codebase. + +import copy + + +def update_config(config: dict, max_individuals: int, device: str): + """Loads the model configuration file for a model, detector and SuperAnimal + + Args: + config: The default model configuration file. + max_individuals: The maximum number of detections to make in an image + device: The device to use to train/run inference on the model + + Returns: + The model configuration for a SuperAnimal-pretrained model. + """ + config = replace_default_values( + config, + num_bodyparts=len(config["metadata"]["bodyparts"]), + num_individuals=max_individuals, + backbone_output_channels=config["model"]["backbone_output_channels"], + ) + config["metadata"]["individuals"] = [f"animal{i}" for i in range(max_individuals)] + + config["device"] = device + if config.get("detector", None) is not None: + config["detector"]["device"] = device + + return config + + +def replace_default_values( + config: dict | list, + num_bodyparts: int | None = None, + num_individuals: int | None = None, + backbone_output_channels: int | None = None, + **kwargs, +) -> dict: + """Replaces placeholder values in a model configuration with their actual values. + + This method allows to create template PyTorch configurations for models with values + such as "num_bodyparts", which are replaced with the number of bodyparts for a + project when making its Pytorch configuration. + + This code can also do some basic arithmetic. You can write "num_bodyparts x 2" (or + any factor other than 2) for location refinement channels, and the number of + channels will be twice the number of bodyparts. You can write + "backbone_output_channels // 2" for the number of channels in a layer, and it will + be half the number of channels output by the backbone. You can write + "num_bodyparts + 1" (such as for DEKR heatmaps, where a "center" bodypart is added). + + The three base placeholder values that can be computed are "num_bodyparts", + "num_individuals" and "backbone_output_channels". You can add more through the + keyword arguments (such as "paf_graph": list[tuple[int, int]] or + "paf_edges_to_keep": list[int] for DLCRNet models). + + Args: + config: the configuration in which to replace default values + num_bodyparts: the number of bodyparts + num_individuals: the number of individuals + backbone_output_channels: the number of backbone output channels + kwargs: other placeholder values to fill in + + Returns: + the configuration with placeholder values replaced + + Raises: + ValueError: if there is a placeholder value who's "updated" value was not + given to the method + """ + + def get_updated_value(variable: str) -> int | list[int]: + var_parts = variable.strip().split(" ") + var_name = var_parts[0] + if updated_values[var_name] is None: + raise ValueError( + f"Found {variable} in the configuration file, but there is no default " + f"value for this variable." + ) + + if len(var_parts) == 1: + return updated_values[var_name] + elif len(var_parts) == 3: + operator, factor = var_parts[1], var_parts[2] + if not factor.isdigit(): + raise ValueError(f"F must be an integer in variable: {variable}") + + factor = int(factor) + if operator == "+": + return updated_values[var_name] + factor + elif operator == "x": + return updated_values[var_name] * factor + elif operator == "//": + return updated_values[var_name] // factor + else: + raise ValueError(f"Unknown operator for variable: {variable}") + + raise ValueError( + f"Found {variable} in the configuration file, but cannot parse it." + ) + + updated_values = { + "num_bodyparts": num_bodyparts, + "num_individuals": num_individuals, + "backbone_output_channels": backbone_output_channels, + **kwargs, + } + + config = copy.deepcopy(config) + if isinstance(config, dict): + keys_to_update = list(config.keys()) + elif isinstance(config, list): + keys_to_update = range(len(config)) + else: + raise ValueError(f"Config to update must be dict or list, found {type(config)}") + + for k in keys_to_update: + if isinstance(config[k], (list, dict)): + config[k] = replace_default_values( + config[k], + num_bodyparts, + num_individuals, + backbone_output_channels, + **kwargs, + ) + elif ( + isinstance(config[k], str) + and config[k].strip().split(" ")[0] in updated_values.keys() + ): + config[k] = get_updated_value(config[k]) + + return config \ No newline at end of file diff --git a/dlclive/modelzoo/utils.py b/dlclive/modelzoo/utils.py new file mode 100644 index 0000000..1aab80b --- /dev/null +++ b/dlclive/modelzoo/utils.py @@ -0,0 +1,175 @@ +""" +Utils for the DLC-Live Model Zoo +""" +# NOTE JR 2026-23-01: This file contains duplicated code from the DeepLabCut main repository. +# This should be removed once a solution is found to address duplicate code. + +import copy +from pathlib import Path +import logging + +from ruamel.yaml import YAML + +from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model +from dlclive.modelzoo.resolve_config import update_config + +_MODELZOO_PATH = Path(__file__).parent + + +def get_super_animal_model_config_path(model_name: str) -> Path: + """Get the path to the model configuration file for a model and validate choice of model""" + cfg_path = _MODELZOO_PATH / 'model_configs' / f"{model_name}.yaml" + if not cfg_path.exists(): + raise FileNotFoundError( + f"Modelzoo model configuration file not found: {cfg_path} " + f"Available models: {list_available_models()}" + ) + return cfg_path + + +def get_super_animal_project_config_path(super_animal: str) -> Path: + """Get the path to the project configuration file for a project and validate choice of project""" + cfg_path = _MODELZOO_PATH / 'project_configs' / f"{super_animal}.yaml" + if not cfg_path.exists(): + raise FileNotFoundError( + f"Modelzoo project configuration file not found: {cfg_path}" + f"Available projects: {list_available_projects()}" + ) + return cfg_path + + +def get_snapshot_folder_path() -> Path: + return _MODELZOO_PATH / 'snapshots' + + +def list_available_models() -> list[str]: + return [p.stem for p in _MODELZOO_PATH.glob('model_configs/*.yaml')] + + +def list_available_projects() -> list[str]: + return [p.stem for p in _MODELZOO_PATH.glob('project_configs/*.yaml')] + + +def list_available_combinations() -> list[str]: + models = list_available_models() + projects = list_available_projects() + combinations = ['_'.join([p, m]) for p in projects for m in models] + return combinations + + +def read_config_as_dict(config_path: str | Path) -> dict: + """ + Args: + config_path: the path to the configuration file to load + + Returns: + The configuration file with pure Python classes + """ + with open(config_path, "r") as f: + cfg = YAML(typ='safe', pure=True).load(f) + + return cfg + + +# NOTE JR 2026-23-01: This is duplicate code, copied from the original DeepLabCut-Live codebase. +def add_metadata(project_config: dict, config: dict,) -> dict: + """Adds metadata to a pytorch pose configuration + + Args: + project_config: the project configuration + config: the pytorch pose configuration + pose_config_path: the path where the pytorch pose configuration will be saved + + Returns: + the configuration with a `meta` key added + """ + config = copy.deepcopy(config) + config["metadata"] = { + "project_path": project_config["project_path"], + "pose_config_path": "", + "bodyparts": project_config.get("multianimalbodyparts") or project_config["bodyparts"], + "unique_bodyparts": project_config.get("uniquebodyparts", []), + "individuals": project_config.get("individuals", ["animal"]), + "with_identity": project_config.get("identity", False), + } + return config + + +# NOTE JR 2026-23-01: This is duplicate code, copied from the original DeepLabCut-Live codebase. +def load_super_animal_config( + super_animal: str, + model_name: str, + detector_name: str | None = None, + max_individuals: int = 30, + device: str | None = None, +) -> dict: + """Loads the model configuration file for a model, detector and SuperAnimal + + Args: + super_animal: The name of the SuperAnimal for which to create the model config. + model_name: The name of the model for which to create the model config. + detector_name: The name of the detector for which to create the model config. + max_individuals: The maximum number of detections to make in an image + device: The device to use to train/run inference on the model + + Returns: + The model configuration for a SuperAnimal-pretrained model. + """ + project_cfg_path = get_super_animal_project_config_path(super_animal=super_animal) + project_config = read_config_as_dict(project_cfg_path) + + model_cfg_path = get_super_animal_model_config_path(model_name=model_name) + model_config = read_config_as_dict(model_cfg_path) + model_config = add_metadata(project_config, model_config) + model_config = update_config(model_config, max_individuals, device) + + if detector_name is None and super_animal != "superanimal_humanbody": + model_config["method"] = "BU" + else: + model_config["method"] = "TD" + if super_animal != "superanimal_humanbody": + detector_cfg_path = get_super_animal_model_config_path( + model_name=detector_name + ) + detector_cfg = read_config_as_dict(detector_cfg_path) + model_config["detector"] = detector_cfg + return model_config + + +def download_super_animal_snapshot(dataset: str, model_name: str) -> Path: + """Downloads a SuperAnimal snapshot + + Args: + dataset: The name of the SuperAnimal dataset for which to download a snapshot. + model_name: The name of the model for which to download a snapshot. + + Returns: + The path to the downloaded snapshot. + + Raises: + RuntimeError if the model fails to download. + """ + snapshot_dir = get_snapshot_folder_path() + model_name = f"{dataset}_{model_name}" + model_filename = f"{model_name}.pt" + model_path = snapshot_dir / model_filename + + if model_path.exists(): + logging.info(f"Snapshot {model_path} already exists, skipping download") + return model_path + + try: + download_huggingface_model( + model_name, target_dir=str(snapshot_dir), rename_mapping=model_filename + ) + + if not model_path.exists(): + raise RuntimeError(f"Failed to download {model_name} to {model_path}") + + except Exception as e: + logging.error(f"Failed to download superanimal snapshot {model_name} to {model_path}: {e}") + raise e + + return model_path + + diff --git a/tests/test_benchmark_script.py b/tests/test_benchmark_script.py index fd7373a..58c1533 100644 --- a/tests/test_benchmark_script.py +++ b/tests/test_benchmark_script.py @@ -4,12 +4,14 @@ from dlclive.engine import Engine -# TODO: JR include separate functional tests for torch and tf backends -@pytest.mark.functional -def test_benchmark_script_runs(tmp_path): +@pytest.fixture +def datafolder(tmp_path): datafolder = tmp_path / "Data-DLC-live-benchmark" download_benchmarking_data(str(datafolder)) + return datafolder +@pytest.mark.functional +def test_benchmark_script_runs_tf_backend(tmp_path, datafolder): dog_models = glob.glob(str(datafolder / "dog" / "*[!avi]")) dog_video = glob.glob(str(datafolder / "dog" / "*.avi"))[0] mouse_models = glob.glob(str(datafolder / "mouse_lick" / "*[!avi]")) @@ -52,3 +54,57 @@ def test_benchmark_script_runs(tmp_path): ) assert any(out_dir.iterdir()) + + +@pytest.mark.parametrize("model_name", ["hrnet_w32", "resnet_50"]) +@pytest.mark.functional +def test_benchmark_script_with_torch_modelzoo(tmp_path, datafolder, model_name): + from dlclive import modelzoo + + # Test configuration + pixels = 4096 # approximately 64x64 pixels, keeping aspect ratio + n_frames = 5 + out_dir = tmp_path / "results" + out_dir.mkdir(exist_ok=True) + + # Export models + model_configs = [ + { + "checkpoint": tmp_path / f"exported_quadruped_{model_name}.pt", + "super_animal": "superanimal_quadruped", + "video_dir": "dog", + }, + { + "checkpoint": tmp_path / f"exported_topviewmouse_{model_name}.pt", + "super_animal": "superanimal_topviewmouse", + "video_dir": "mouse_lick", + }, + ] + + for config in model_configs: + modelzoo.export_modelzoo_model( + export_path=config["checkpoint"], + super_animal=config["super_animal"], + model_name=model_name, + ) + assert config["checkpoint"].exists(), f"Failed to export {config['super_animal']} model" + assert config["checkpoint"].stat().st_size > 0, f"Exported {config['super_animal']} model is empty" + + # Get video paths and run benchmarks + for config in model_configs: + video_dir = datafolder / config["video_dir"] + video_path = list(video_dir.glob("*.avi"))[0] + print(f"Running {config['checkpoint'].stem}") + benchmark_videos( + model_path=config["checkpoint"], + model_type="pytorch", + video_path=video_path, + output=str(out_dir), + n_frames=n_frames, + pixels=pixels, + ) + + # Assertions: verify output files were created + output_files = list(out_dir.iterdir()) + assert len(output_files) > 0, "No output files were created by benchmark_videos" + assert any(f.suffix == ".pickle" for f in output_files), "No pickle files found in output directory" \ No newline at end of file diff --git a/tests/test_modelzoo.py b/tests/test_modelzoo.py new file mode 100644 index 0000000..1a7e6db --- /dev/null +++ b/tests/test_modelzoo.py @@ -0,0 +1,51 @@ +# NOTE JR 2026-23-01: This is duplicate code, copied from the original DeepLabCut-Live codebase. + +import os + +import pytest +import dlclibrary +from dlclibrary.dlcmodelzoo.modelzoo_download import MODELOPTIONS + +from dlclive import modelzoo + + +@pytest.mark.parametrize( + "super_animal", ["superanimal_quadruped", "superanimal_topviewmouse"] +) +@pytest.mark.parametrize("model_name", ["hrnet_w32"]) +@pytest.mark.parametrize("detector_name", [None, "fasterrcnn_resnet50_fpn_v2"]) +def test_get_config_model_paths(super_animal, model_name, detector_name): + model_config = modelzoo.load_super_animal_config( + super_animal=super_animal, + model_name=model_name, + detector_name=detector_name, + ) + + assert isinstance(model_config, dict) + if detector_name is None: + assert model_config["method"].lower() == "bu" + assert "detector" not in model_config + else: + assert model_config["method"].lower() == "td" + assert "detector" in model_config + + +def test_download_huggingface_model(tmp_path_factory, model="full_cat"): + folder = tmp_path_factory.mktemp("temp") + dlclibrary.download_huggingface_model(model, str(folder)) + + assert os.path.exists(folder / "pose_cfg.yaml") + assert any(f.startswith("snapshot-") for f in os.listdir(folder)) + # Verify that the Hugging Face folder was removed + assert not any(f.startswith("models--") for f in os.listdir(folder)) + + +def test_download_huggingface_wrong_model(): + with pytest.raises(ValueError): + dlclibrary.download_huggingface_model("wrong_model_name") + + +@pytest.mark.skip(reason="slow") +@pytest.mark.parametrize("model", MODELOPTIONS) +def test_download_all_models(tmp_path_factory, model): + test_download_huggingface_model(tmp_path_factory, model) \ No newline at end of file