How to use GPU for inference

[Unkrible]

Hello, thanks for the open-source of awesome work!
However, I have the question that as guided in the Inference by HumanHeadSegmentationPipeline .predict, GPU can not be used to accelerate the inference process.
Thus, I suggestion to add self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') in class HumanHeadSegmentationPipeline, and then modify the function predict(self, image) as mdl_out = self._model(preprocessed_image.to(self.device)).detach().cpu()

[wiktorlazarski]

Feel free to make a PR with all proposed changes 😉.

[9527-csroad]

I find this problem and test it. After two images, cpu inference time around 2.1s while gpu inference time around 1.4s.
Now, we can set the device like this:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
visualizer = vis.VisualizationModule()
segmentation_pipeline = seg_pipeline.HumanHeadSegmentationPipeline(device=device)

if you not set the device, it will use cpu to inference by default.
And I also find the visualization will consume one part time. If you are strict with time, just save segmented_region like this:

segmentation_map = segmentation_pipeline.predict(image)
segmented_region = image * cv2.cvtColor(segmentation_map, cv2.COLOR_GRAY2RGB)
pil_image = Image.fromarray(segmented_region)
pil_image.save(save_path)

I test it with cpu and it only consume around 0.8s for a single image. If you use gpu, it only take around 0.15s for a single image.

[wiktorlazarski]

Hey @9527-csroad,

thanks for having a look at it 🙌. Amazing work 🚀! I love a time benchmarking. Also, huge shoutout for the PR #30. I promise I'll have a look at it later today.