Merge pull request DefTruth#57 from DefTruth/dev

fix(pipnet): fixed cuda error in pipnet. (DefTruth#56)

README.md

@@ -82,8 +82,8 @@
 ## 🛠️Installation
 you can install **torchlm** directly from [pypi](https://pypi.org/project/torchlm/). 
 ```shell
-pip install torchlm>=0.1.6.9 # or install the latest pypi version `pip install torchlm`
-pip install torchlm>=0.1.6.9 -i https://pypi.org/simple/ # or install from specific pypi mirrors use '-i'
+pip install torchlm>=0.1.6.10 # or install the latest pypi version `pip install torchlm`
+pip install torchlm>=0.1.6.10 -i https://pypi.org/simple/ # or install from specific pypi mirrors use '-i'
 ```
 or install from source if you want the latest torchlm and install it in editable mode with `-e`.
 ```shell
@@ -400,16 +400,18 @@ import torchlm
 from torchlm.tools import faceboxesv2
 from torchlm.models import pipnet
 
-torchlm.runtime.bind(faceboxesv2())
+torchlm.runtime.bind(faceboxesv2(device="cpu"))  # set device="cuda" if you want to run with CUDA
+# set map_location="cuda" if you want to run with CUDA
 torchlm.runtime.bind(
   pipnet(backbone="resnet18", pretrained=True,  
          num_nb=10, num_lms=98, net_stride=32, input_size=256,
-         meanface_type="wflw", map_location="cpu", checkpoint=None)
+         meanface_type="wflw", map_location="cpu", checkpoint=None) 
 ) # will auto download pretrained weights from latest release if pretrained=True
 landmarks, bboxes = torchlm.runtime.forward(image)
 image = torchlm.utils.draw_bboxes(image, bboxes=bboxes)
 image = torchlm.utils.draw_landmarks(image, landmarks=landmarks)
 ```
+
 #### Inference on ONNXRuntime Backend
 ```python
 import torchlm

test/pipnet.py

@@ -164,8 +164,8 @@ def test_pipnet_runtime_ort():
 
 if __name__ == "__main__":
     test_pipnet_runtime()
-    test_pipnet_training()
-    test_pipnet_evaluating()
-    test_pipnet_exporting()
-    test_pipnet_meanface()
-    test_pipnet_runtime_ort()
+    # test_pipnet_training()
+    # test_pipnet_evaluating()
+    # test_pipnet_exporting()
+    # test_pipnet_meanface()
+    # test_pipnet_runtime_ort()

torchlm/__init__.py

@@ -1,5 +1,5 @@
 # Versions
-__version__ = '0.1.6.9'
+__version__ = '0.1.6.10'
 # Transforms Module: 100+ transforms available, can bind torchvision and
 # albumentations into torchlm pipeline with autodtype wrapper.
 from .transforms import *

torchlm/models/pipnet/_impls.py

@@ -285,11 +285,12 @@ def _detecting_impl(
         )
 
     net.eval()
+    device = next(net.parameters()).device
 
     height, width, _ = image.shape
     image: np.ndarray = cv2.resize(image, (net.input_size, net.input_size))  # 256, 256
     image: Tensor = torch.from_numpy(_normalize(img=image)).contiguous().unsqueeze(0)  # (1,3,256,256)
-    outputs_cls, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y = net.forward(image)
+    outputs_cls, outputs_x, outputs_y, outputs_nb_x, outputs_nb_y = net.forward(image.to(device))
     # (1,68,8,8)
     tmp_batch, tmp_channel, tmp_height, tmp_width = outputs_cls.size()
     assert tmp_batch == 1