From 4acb0c887d7e6f0c9fb638edbb10e31f5377b808 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 2 Jun 2024 17:42:46 -0400
Subject: [PATCH] update dwpose installer

---
 install-mm.py                           | 44 +++++++++++++++++++++++++
 modules/control/proc/dwpose/__init__.py | 37 +++++++++++++--------
 2 files changed, 68 insertions(+), 13 deletions(-)
 create mode 100644 install-mm.py

diff --git a/install-mm.py b/install-mm.py
new file mode 100644
index 000000000..941f01e52
--- /dev/null
+++ b/install-mm.py
@@ -0,0 +1,44 @@
+import os
+from installer import setup_logging
+setup_logging()
+
+
+checked_ok = False
+
+
+def check_dependencies():
+    from installer import installed, pip, log
+    global checked_ok # pylint: disable=global-statement
+    debug = log.trace if os.environ.get('SD_DWPOSE_DEBUG', None) is not None else lambda *args, **kwargs: None
+    packages = [
+        'openmim==0.3.9',
+        'mmengine==0.10.4',
+        'mmcv==2.1.0',
+        'mmpose==1.3.1',
+        'mmdet==3.3.0',
+    ]
+    status = [installed(p, reload=False, quiet=False) for p in packages]
+    status.append(False)
+    debug(f'DWPose required={packages} status={status}')
+    if not all(status):
+        log.info(f'Installing DWPose dependencies: {[packages]}')
+        cmd = 'install --upgrade --no-deps --force-reinstall '
+        pkgs = ' '.join(packages)
+        res = pip(cmd + pkgs, ignore=False, quiet=False)
+        debug(f'DWPose pip install: {res}')
+    try:
+        import pkg_resources
+        import imp # pylint: disable=deprecated-module
+        imp.reload(pkg_resources)
+        import mmcv # pylint: disable=unused-import
+        import mmengine # pylint: disable=unused-import
+        import mmpose # pylint: disable=unused-import
+        import mmdet # pylint: disable=unused-import
+        debug('DWPose import ok')
+        checked_ok = True
+    except Exception as e:
+        log.error(f'DWPose: {e}')
+    return checked_ok
+
+
+check_dependencies()
diff --git a/modules/control/proc/dwpose/__init__.py b/modules/control/proc/dwpose/__init__.py
index 4ae62c133..4469e10c8 100644
--- a/modules/control/proc/dwpose/__init__.py
+++ b/modules/control/proc/dwpose/__init__.py
@@ -9,7 +9,7 @@ os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
 import cv2
 import numpy as np
 from PIL import Image
-from installer import installed, install, log
+from installer import installed, pip, log
 from modules.control.util import HWC3, resize_image
 from .draw import draw_bodypose, draw_handpose, draw_facepose
 checked_ok = False
@@ -17,24 +17,35 @@ checked_ok = False
 
 def check_dependencies():
     global checked_ok # pylint: disable=global-statement
+    debug = log.trace if os.environ.get('SD_DWPOSE_DEBUG', None) is not None else lambda *args, **kwargs: None
     packages = [
-        ('openmim==0.3.9', 'openmim'),
-        ('mmengine==0.10.4', 'mmengine'),
-        ('mmcv==2.1.0', 'mmcv'),
-        ('mmpose==1.3.1', 'mmpose'),
-        ('mmdet==3.3.0', 'mmdet'),
+        'openmim==0.3.9',
+        'mmengine==0.10.4',
+        'mmcv==2.1.0',
+        'mmpose==1.3.1',
+        'mmdet==3.3.0',
     ]
-    packages = []
-    for pkg in packages:
-        if not installed(pkg[1], reload=True, quiet=True):
-            install(pkg[0], pkg[1], ignore=False, no_deps=True)
+    status = [installed(p, reload=False, quiet=False) for p in packages]
+    debug(f'DWPose required={packages} status={status}')
+    if not all(status):
+        log.info(f'Installing DWPose dependencies: {[packages]}')
+        cmd = 'install --upgrade --no-deps --force-reinstall '
+        pkgs = ' '.join(packages)
+        res = pip(cmd + pkgs, ignore=False, quiet=False)
+        debug(f'DWPose pip install: {res}')
     try:
+        import pkg_resources
+        import imp # pylint: disable=deprecated-module
+        imp.reload(pkg_resources)
         import mmcv # pylint: disable=unused-import
+        import mmengine # pylint: disable=unused-import
+        import mmpose # pylint: disable=unused-import
+        import mmdet # pylint: disable=unused-import
+        debug('DWPose import ok')
         checked_ok = True
-        return True
     except Exception as e:
         log.error(f'DWPose: {e}')
-        return False
+    return checked_ok
 
 
 def draw_pose(pose, H, W):
@@ -48,9 +59,9 @@ def draw_pose(pose, H, W):
     canvas = draw_bodypose(canvas, candidate, subset)
     canvas = draw_handpose(canvas, hands)
     canvas = draw_facepose(canvas, faces)
-
     return canvas
 
+
 class DWposeDetector:
     def __init__(self, det_config=None, det_ckpt=None, pose_config=None, pose_ckpt=None, device="cpu"):
         self.pose_estimation = None