small changes

Fannovel16 · Dec 7, 2023 · 4b2349c · 4b2349c
1 parent 4f9ac2f
commit 4b2349c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/src/controlnet_aux/densepose/__init__.py b/src/controlnet_aux/densepose/__init__.py
@@ -1,7 +1,4 @@
-import functools
-import os
-import warnings
-import torchvision #lack of nmp
+import torchvision # Fix issue Unknown builtin op: torchvision::nms
 import cv2
 import numpy as np
 import torch
@@ -10,12 +7,21 @@
 from PIL import Image
 
 from controlnet_aux.util import HWC3, resize_image_with_pad, common_input_validate, annotator_ckpts_path, custom_hf_download
-from .densepose import civitai_vis, densepose_chart_predictor_output_to_result_with_confidences
+from .densepose import DensePoseMaskedColormapResultsVisualizer, _extract_i_from_iuvarr, densepose_chart_predictor_output_to_result_with_confidences
+
+N_PART_LABELS = 24
 
 class DenseposeDetector:
     def __init__(self, model):
         self.dense_pose_estimation = model
         self.device = "cpu"
+        self.result_visualizer = DensePoseMaskedColormapResultsVisualizer(
+            cmap=cv2.COLORMAP_PARULA, 
+            alpha=1, 
+            data_extractor=_extract_i_from_iuvarr, 
+            segm_extractor=_extract_i_from_iuvarr, 
+            val_scale = 255.0 / N_PART_LABELS
+        )
 
     @classmethod
     def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=annotator_ckpts_path):
@@ -43,7 +49,7 @@ def __call__(self, input_image, detect_resolution=512, output_type="pil", upscal
         extractor = densepose_chart_predictor_output_to_result_with_confidences
         densepose_results = [extractor(pred_boxes[i:i+1], corase_segm[i:i+1], fine_segm[i:i+1], u[i:i+1], v[i:i+1]) for i in range(len(pred_boxes))]
 
-        hint_image = civitai_vis.visualize(hint_image_canvas, densepose_results)
+        hint_image = self.result_visualizer.visualize(hint_image_canvas, densepose_results)
         hint_image = cv2.cvtColor(hint_image, cv2.COLOR_BGR2RGB)
         hint_image = torch.from_numpy(hint_image)
 

diff --git a/src/controlnet_aux/densepose/densepose.py b/src/controlnet_aux/densepose/densepose.py
@@ -10,7 +10,6 @@
 
 Image = np.ndarray
 Boxes = torch.Tensor
-N_PART_LABELS = 24
 ImageSizeType = Tuple[int, int]
 _RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray]
 IntTupleBox = Tuple[int, int, int, int]
@@ -346,4 +345,3 @@ def resample_uv_tensors_to_bbox(
         uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
     return uv
 
-civitai_vis = DensePoseMaskedColormapResultsVisualizer(cmap=cv2.COLORMAP_PARULA, alpha=1, data_extractor=_extract_i_from_iuvarr, segm_extractor=_extract_i_from_iuvarr, val_scale = 255.0 / N_PART_LABELS)