Merge pull request #1 from AdaptiveMotorControlLab/MMathisLab/fmpose3d

Update project name; fix env errors;

Author: xiu-cs

README.md

@@ -22,8 +22,8 @@
 from PIL import Image
 import matplotlib.gridspec as gridspec
 import imageio
-from fmpose.animals.common.arguments import opts as parse_args
-from fmpose.common.camera import normalize_screen_coordinates, camera_to_world
+from fmpose3d.animals.common.arguments import opts as parse_args
+from fmpose3d.common.camera import normalize_screen_coordinates, camera_to_world
 
 sys.path.append(os.getcwd())
 
@@ -46,7 +46,7 @@
     CFM = getattr(module, "Model")
 else:
     # Load model from installed fmpose package
-    from fmpose.models import Model as CFM
+    from fmpose3d.models import Model as CFM
 
 from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images
 

animals/README.md

@@ -7,7 +7,7 @@ sh_file='vis_animals.sh'
 # n_joints=26
 # out_joints=26
 
-model_path='../../fmpose/animals/models/model_animal3d.py'
+model_path='../pre_trained_models/animal3d_pretrained_weights/model_animal3d.py'
 saved_model_path='../pre_trained_models/animal3d_pretrained_weights/CFM_154_4403_best.pth'
 
 # path='./images/image_00068.jpg'  # single image

animals/demo/vis_animals.py

@@ -15,7 +15,7 @@
 from einops import rearrange
 from timm.models.layers import DropPath
 
-from fmpose.animals.models.graph_frames import Graph
+from fmpose3d.animals.models.graph_frames import Graph
 
 
 class TimeEmbedding(nn.Module):
@@ -148,7 +148,7 @@ def forward(self, x):
         x = res2 + self.drop_path(x)
         return x
 
-class FMPose(nn.Module):
+class FMPose3D(nn.Module):
     def __init__(self, depth, embed_dim, channels_dim, tokens_dim, adj, drop_rate=0.10, length=27):
         super().__init__()
         drop_path_rate = 0.2
@@ -220,7 +220,7 @@ def __init__(self, args):
         self.encoder_pose_2d = encoder(2, args.channel//2, args.channel//2-self.t_embed_dim//2)
         self.encoder_y_t = encoder(3, args.channel//2, args.channel//2-self.t_embed_dim//2)
 
-        self.FMPose = FMPose(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256
+        self.FMPose3D = FMPose3D(args.layers, args.channel, args.d_hid, args.token_dim, self.A, length=args.n_joints) # 256
         self.pred_mu = decoder(args.channel, args.channel//2, 3)
 
     def forward(self, pose_2d, y_t, t):
@@ -239,7 +239,7 @@ def forward(self, pose_2d, y_t, t):
         in_emb = rearrange(in_emb, 'b f j c -> (b f) j c').contiguous() # (B*F,J,in)
 
         # encoder -> model -> regression head
-        h = self.FMPose(in_emb)
+        h = self.FMPose3D(in_emb)
         v = self.pred_mu(h)                                  # (B*F,J,3)
 
         v = rearrange(v, '(b f) j c -> b f j c', b=b, f=f).contiguous() # (B,F,J,3)

animals/demo/vis_animals.sh

@@ -15,9 +15,9 @@
 import numpy as np
 from tqdm import tqdm
 import torch.optim as optim
-from fmpose.animals.common.arguments import opts as parse_args
-from fmpose.animals.common.utils import *
-from fmpose.animals.common.animal3d_dataset import TrainDataset
+from fmpose3d.animals.common.arguments import opts as parse_args
+from fmpose3d.animals.common.utils import *
+from fmpose3d.animals.common.animal3d_dataset import TrainDataset
 import time
 
 args = parse_args().parse()
@@ -39,7 +39,7 @@
     CFM = getattr(module, "Model")
 else:
     # Load model from installed fmpose package
-    from fmpose.animals.models import Model as CFM
+    from fmpose3d.animals.models import Model as CFM
 
 def train(opt, actions, train_loader, model, optimizer, epoch):
     return step('train', opt, actions, train_loader, model, optimizer, epoch)

animals/models/model_animals.py

@@ -11,7 +11,7 @@ n_joints=26
 out_joints=26
 epochs=300
 # model_path='models/model_animals.py'
-model_path="" # when the path is empty, the model will be loaded from the installed fmpose package
+model_path='./pre_trained_models/animal3d_pretrained_weights/model_animal3d.py' # when the path is empty, the model will be loaded from the installed fmpose package
 saved_model_path='./pre_trained_models/animal3d_pretrained_weights/CFM_154_4403_best.pth'
 # root path denotes the path to the original dataset
 root_path="./dataset/"

animals/scripts/main_animal3d.py

@@ -0,0 +1 @@
+Download the preprocessed datasets from [here](https://drive.google.com/drive/folders/112GPdRC9IEcwcJRyrLJeYw9_YV4wLdKC?usp=sharing) and place them in this folder.

animals/scripts/test_animal3d.sh

@@ -19,12 +19,12 @@
 sys.path.append(os.getcwd())
 
 # Auto-download checkpoint files if missing
-from fmpose.lib.checkpoint.download_checkpoints import ensure_checkpoints
+from fmpose3d.lib.checkpoint.download_checkpoints import ensure_checkpoints
 ensure_checkpoints()
 
-from fmpose.lib.preprocess import h36m_coco_format, revise_kpts
-from fmpose.lib.hrnet.gen_kpts import gen_video_kpts as hrnet_pose
-from fmpose.common.arguments import opts as parse_args
+from fmpose3d.lib.preprocess import h36m_coco_format, revise_kpts
+from fmpose3d.lib.hrnet.gen_kpts import gen_video_kpts as hrnet_pose
+from fmpose3d.common.arguments import opts as parse_args
 
 args = parse_args().parse()
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
@@ -39,7 +39,7 @@
     spec.loader.exec_module(module)
     CFM = getattr(module, 'Model')
 
-from fmpose.common.camera import *
+from fmpose3d.common.camera import *
 
 import matplotlib
 import matplotlib.pyplot as plt