remove frozen requirements (#105)

* clean up requirements

* fix typechecking errors introduced by unfreezing requirements

* replace skimage op with cv2

Author: themattinthehatt

README.md

@@ -57,11 +57,15 @@ cd <SOME_FOLDER>
 git clone https://github.com/danbider/lightning-pose.git
 ```
 
-Then move into the newly-created repository folder, and install dependencies:
+Then move into the newly-created repository folder:
 ```console
 cd lightning-pose
-pip install -e .
 ```
+and install dependencies using one of the lines below that suits your needs best:
+* `pip install -e . `: basic installation, covers most use-cases (note the period!)
+* `pip install -e .[dev] `: basic install + dev tools
+* `pip install -e .[extra_models] `: basic install + tools for loading resnet-50 simclr weights
+* `pip install -e .[dev,extra_models] `: install all available requirements
 
 This installation might take between 3-10 minutes, depending on your machine and internet connection.
 

docs/contributing.md

@@ -6,9 +6,18 @@ If you have found a bug or would like to request a minor change, please
 
 In order to contribute code to the repo, please follow the steps below.
 
-You will also need to install the following dev packages:
+Whenever you initially install the lightning pose repo, instead of
 ```bash
-pip install flake8 isort
+pip install -e .
+```
+run
+```bash
+pip install -e .[dev]
+```
+
+Alternatively, if you have already installed the repo, install the following dev packages:
+```bash
+pip install black flake8 isort
 ```
 
 ### Create a pull request

lightning_pose/models/__init__.py

@@ -6,7 +6,7 @@
     "resnet50",
     "resnet101",
     "resnet152",
-    "resnet50_contrastive",
+    "resnet50_contrastive",  # needs extra install: pip install -e .[extra_models]
     "resnet50_animal_apose",
     "resnet50_animal_ap10k",
     "resnet50_human_jhmdb",
@@ -15,6 +15,6 @@
     "efficientnet_b0",
     "efficientnet_b1",
     "efficientnet_b2",
-    "vit_h_sam",
-    "vit_b_sam",
+    # "vit_h_sam",
+    # "vit_b_sam",
 ]

lightning_pose/models/backbones/torchvision.py

@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from typing import Tuple
 
 import torch
 import torchvision.models as tvmodels
@@ -28,8 +29,14 @@ def build_backbone(
 
     if backbone_arch == "resnet50_contrastive":
         # load resnet50 pretrained using SimCLR on imagenet
-        from pl_bolts.models.self_supervised import SimCLR
-
+        try:
+            from pl_bolts.models.self_supervised import SimCLR
+        except ImportError:
+            raise Exception(
+                "lightning-bolts package is not installed.\n"
+                "Run `pip install lightning-bolts` "
+                "in order to access 'resnet50_contrastive' backbone"
+            )
         ckpt_url = "https://pl-bolts-weights.s3.us-east-2.amazonaws.com/simclr/bolts_simclr_imagenet/simclr_imagenet.ckpt"  # noqa: E501
         simclr = SimCLR.load_from_checkpoint(ckpt_url, strict=False)
         base = simclr.encoder

lightning_pose/models/heatmap_tracker.py

@@ -1,6 +1,6 @@
 """Models that produce heatmaps of keypoints from images."""
 
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 from kornia.filters import filter2d
@@ -9,6 +9,7 @@
 from omegaconf import DictConfig
 from torch import nn
 from torchtyping import TensorType
+from typeguard import typechecked
 from typing_extensions import Literal
 
 from lightning_pose.data.utils import (
@@ -349,6 +350,7 @@ def predict_step(
             return predicted_keypoints, confidence
 
 
+@typechecked
 class SemiSupervisedHeatmapTracker(SemiSupervisedTrackerMixin, HeatmapTracker):
     """Model produces heatmaps of keypoints from labeled/unlabeled images."""
 
@@ -411,7 +413,7 @@ def __init__(
         # self.register_buffer("total_unsupervised_importance", torch.tensor(1.0))
         self.total_unsupervised_importance = torch.tensor(1.0)
 
-    def get_loss_inputs_unlabeled(self, batch: UnlabeledBatchDict) -> dict:
+    def get_loss_inputs_unlabeled(self, batch: UnlabeledBatchDict) -> Dict:
         """Return predicted heatmaps and their softmaxes (estimated keypoints)."""
         # images -> heatmaps
         predicted_heatmaps = self.forward(batch["frames"])

lightning_pose/models/heatmap_tracker_mhcrnn.py

@@ -1,6 +1,6 @@
 """Models that produce heatmaps of keypoints from images."""
 
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 from kornia.geometry.subpix import spatial_softmax2d
@@ -259,7 +259,7 @@ def __init__(
         # self.register_buffer("total_unsupervised_importance", torch.tensor(1.0))
         self.total_unsupervised_importance = torch.tensor(1.0)
 
-    def get_loss_inputs_unlabeled(self, batch: UnlabeledBatchDict) -> dict:
+    def get_loss_inputs_unlabeled(self, batch: UnlabeledBatchDict) -> Dict:
         """Return predicted heatmaps and their softmaxes (estimated keypoints)."""
         # images -> heatmaps
         pred_heatmaps_crnn, pred_heatmaps_sf = self.forward(batch["frames"])

lightning_pose/models/regression_tracker.py

@@ -1,6 +1,6 @@
 """Models that produce (x, y) coordinates of keypoints from images."""
 
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 from omegaconf import DictConfig
@@ -204,7 +204,7 @@ def __init__(
         self.total_unsupervised_importance = torch.tensor(1.0)
         # self.register_buffer("total_unsupervised_importance", torch.tensor(1.0))
 
-    def get_loss_inputs_unlabeled(self, batch: UnlabeledBatchDict) -> dict:
+    def get_loss_inputs_unlabeled(self, batch: UnlabeledBatchDict) -> Dict:
         """Return predicted heatmaps and their softmaxes (estimated keypoints)."""
         predicted_keypoints = self.forward(batch["frames"])
         # undo augmentation if needed

lightning_pose/utils/io.py

@@ -190,7 +190,7 @@ def get_videos_in_dir(video_dir: str, return_mp4_only: bool = True) -> List[str]
 
 
 @typechecked
-def check_video_paths(video_paths: Union[List[str], str]) -> list:
+def check_video_paths(video_paths: Union[List[str], str]) -> List[str]:
     # get input data
     if isinstance(video_paths, list):
         # presumably a list of files
@@ -203,8 +203,7 @@ def check_video_paths(video_paths: Union[List[str], str]) -> list:
         filenames = get_videos_in_dir(video_paths)
     else:
         raise ValueError(
-            "`video_paths_list` must be a list of files, a single file, "
-            + "or a directory name"
+            "`video_paths_list` must be a list of files, a single file, or a directory name"
         )
     for filename in filenames:
         assert filename.endswith(".mp4"), "video files must be mp4 format!"

lightning_pose/utils/predictions.py

@@ -4,14 +4,14 @@
 import time
 from typing import List, Optional, Tuple, Union
 
+import cv2
 import lightning.pytorch as pl
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import torch
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import LightningModule
-from skimage.draw import disk
 from torchtyping import TensorType
 from tqdm import tqdm
 from typeguard import typechecked
@@ -690,7 +690,7 @@ def create_labeled_video(
     xs_arr,
     ys_arr,
     mask_array=None,
-    dotsize=5,
+    dotsize=4,
     colormap="cool",
     fps=None,
     filename="movie.mp4",
@@ -719,14 +719,10 @@ def create_labeled_video(
     colors = make_cmap(n_keypoints, cmap=colormap)
 
     nx, ny = clip.size
-    duration = int(clip.duration - clip.start)
+    dur = int(clip.duration - clip.start)
     fps_og = clip.fps
 
-    print(
-        "Duration of video [s]: {}, recorded with {} fps!".format(
-            np.round(duration, 2), np.round(fps_og, 2)
-        )
-    )
+    print(f"Duration of video [s]: {np.round(dur, 2)}, recorded at {np.round(fps_og, 2)} fps!")
 
     # add marker to each frame t, where t is in sec
     def add_marker(get_frame, t):
@@ -742,8 +738,13 @@ def add_marker(get_frame, t):
             if mask_array[index, bpindex]:
                 xc = min(int(xs_arr[index, bpindex]), nx - 1)
                 yc = min(int(ys_arr[index, bpindex]), ny - 1)
-                rr, cc = disk(center=(yc, xc), radius=dotsize, shape=(ny, nx))
-                frame[rr, cc, :] = colors[bpindex]
+                frame = cv2.circle(
+                    frame,
+                    center=(xc, yc),
+                    radius=dotsize,
+                    color=colors[bpindex].tolist(),
+                    thickness=-1
+                )
         return frame
 
     clip_marked = clip.fl(add_marker)

lightning_pose/utils/scripts.py

@@ -1,7 +1,7 @@
 """Helper functions to build pipeline components from config dictionary."""
 
 import os
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import imgaug.augmenters as iaa
 import lightning.pytorch as pl