Merge pull request #152 from deruyter92/jaap/add_modelzoo_functionality

functionality for direct downloading and exporting the pytorch models from the modelzoo

Author: deruyter92

.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/)

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

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

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

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

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