realtime pose stuff in blender with opencv

opencv_pose_func.py

# based upon https://github.com/jkirsons/FacialMotionCapture_v2/blob/master/OpenCVAnimOperator.py
# with elements taken from https://github.com/legolas123/cv-tricks.com/blob/master/OpenCV/Pose_Estimation/run_pose.py

# in blender, run this script second

import bpy
import sys
import time
import numpy

# For custom opencv-python, set our system path to read from our virtualenv

venv_path = '/home/zumid/.pyenv/versions/opencv-cuda/lib/python3.8/site-packages'

if venv_path not in sys.path:
    sys.path.insert(0, venv_path)

# ...then import our custom CUDA-enabled CV2
import cv2
    
class OpenCVAnimOperator(bpy.types.Operator):
    """Operator which runs its self from a timer"""
    bl_idname = "wm.opencv_operator"
    bl_label = "OpenCV Animation Operator"
    
    # models for pose estimation from OpenPose
    # while I was supposed to download these from http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/body_25/,
    # the server seems to be down
    # so refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/issues/1602
    # and use link https://drive.google.com/file/d/1QCSxJZpnWvM00hx49CJ2zky7PWGzpcEh/view to get the models
    # (warning, 2.6gb file)
    pose_txt = "/run/media/zumid/Other/openpose_training_data/pose/body_25/pose_deploy.prototxt"
    pose_model = "/run/media/zumid/Other/openpose_training_data/pose/body_25/pose_iter_584000.caffemodel"
    
    NEEDED_BODY_PARTS = {
    "LShoulder" : 5,
    "LElbow"    : 6,
    "LWrist"    : 7,
    "RShoulder" : 2,
    "RElbow"    : 3,
    "RWrist"    : 4,
    "Nose"      : 0,
    "Neck"      : 1
    }

    _timer = None
    _webcam  = None
    
    threshold = 0.2
    
    show_cam = False
    conceal_cam = False
    
    smooth_param = 3 # how many frames to smooth out
    
    webcam_size = (640, 480) # make sure the widthxheight mode is supported by webcam
    target_size = (320, 240) # input size for the neural network to work on, smaller = faster
    
    # load the pose models
    pose_net = cv2.dnn.readNetFromCaffe(pose_txt, pose_model)
    pose_net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
    pose_net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

    stop :bpy.props.BoolProperty()
    
    # Keeps a moving average of given length
    def smooth_value(self, name, length, value):
        if not hasattr(self, 'smooth'):
            self.smooth = {}
        if not name in self.smooth:
            self.smooth[name] = numpy.array([value])
        else:
            self.smooth[name] = numpy.insert(arr=self.smooth[name], obj=0, values=value)
            if self.smooth[name].size > length:
                self.smooth[name] = numpy.delete(self.smooth[name], self.smooth[name].size-1, 0)
        sum = 0
        for val in self.smooth[name]:
            sum += val
        return sum / self.smooth[name].size

    def modal(self, context, event):
        if (event.type in {'ESC'}) or self.stop == True:
            self.cancel(context)
            return {'CANCELLED'}

        if event.type == 'TIMER':
            self.init_camera()
            # Loop here
            
            _, img = self._webcam.read()
            
            img_width = img.shape[1]
            img_height = img.shape[0]
            
            # generate a 4D blob for processing
            in_blob = cv2.dnn.blobFromImage(img, 1.0/255,
                        self.target_size, (0,0,0),
                        swapRB=False, crop=False)
            
            self.pose_net.setInput(in_blob)
            
            # process the image
            start_t = time.time()
            out = self.pose_net.forward()
            print("Calc time: {}".format(time.time() - start_t))
            
            out_width = out.shape[3]
            out_height = out.shape[2]
            
            if self.conceal_cam:
                for i in range(len(img)):
                        # conceal my picture
                        img[i] = 0
            
            for part in self.NEEDED_BODY_PARTS:
                # confidence map of the body part
                prob_map = out[0, self.NEEDED_BODY_PARTS[part], :, :]
                
                # global maxima
                min, prob, min, point = cv2.minMaxLoc(prob_map)
                
                # scale the point
                x = point[0] * (img_width / out_width)
                y = point[1] * (img_height / out_height)
                
                if prob > self.threshold:
                    if self.show_cam:
                        # Visualize our point in the webcam
                        cv2.circle(img,
                                    (int(x), int(y)),
                                    4,
                                    (0,255,255),
                                    thickness=-1,
                                    lineType=cv2.FILLED
                                    )
                        cv2.putText(img,
                                    "{}".format(part),
                                    (int(x), int(y)),
                                    cv2.FONT_HERSHEY_SIMPLEX,
                                    .7,
                                    (0,0,255),
                                    1,
                                    lineType=cv2.LINE_AA
                                    )
                    # Affect the 3D space
                    # only objects which are titled "CAPITAL_TARGET"
                    # my setup is basically empties which then can control the armature bones
                    if part.upper()+"_TARGET" in bpy.data.objects:
                        target_obj = bpy.data.objects[part.upper()+"_TARGET"]
                        
                        # 2d(x) -> 3d(x)
                        target_obj.location[0] = \
                        self.smooth_value(part+"_x", self.smooth_param, point[0] * 0.4)
                        
                        # 2d(y) -> 3d(z)
                        target_obj.location[2] = \
                        self.smooth_value(part+"_y", self.smooth_param, (-point[1]+28) * 0.48)
            if self.show_cam:
                cv2.imshow("Output", img)
            cv2.waitKey(1)

        return {'PASS_THROUGH'}
    
    def init_camera(self):
        if self._webcam == None:
            self._webcam = cv2.VideoCapture(0)
            self._webcam.set(cv2.CAP_PROP_FRAME_WIDTH, self.webcam_size[0])
            self._webcam.set(cv2.CAP_PROP_FRAME_HEIGHT, self.webcam_size[1])
            self._webcam.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            time.sleep(0.1)
    
    def stop_playback(self, scene):
        print(format(scene.frame_current) + " / " + format(scene.frame_end))
        if scene.frame_current == scene.frame_end:
            bpy.ops.screen.animation_cancel(restore_frame=False)
        
    def execute(self, context):
        bpy.app.handlers.frame_change_pre.append(self.stop_playback)

        wm = context.window_manager
        self._timer = wm.event_timer_add(0.02, window=context.window)
        wm.modal_handler_add(self)
        return {'RUNNING_MODAL'}

    def cancel(self, context):
        wm = context.window_manager
        wm.event_timer_remove(self._timer)
        cv2.destroyAllWindows()
        self._webcam.release()
        self._webcam = None

def register():
    bpy.utils.register_class(OpenCVAnimOperator)

def unregister():
    bpy.utils.unregister_class(OpenCVAnimOperator)

if __name__ == "__main__":
    register()

opencv_pose_panel.py

# based on https://github.com/jkirsons/FacialMotionCapture_v2/blob/master/OpenCVAnim.py

# in blender, run this script first

import bpy

class OBJECT_MT_OpenCVPanel(bpy.types.WorkSpaceTool):
    """Creates a Panel in the Object properties window"""
    bl_label = "OpenCV Animation"
    bl_space_type = 'VIEW_3D'
    bl_context_mode='OBJECT'
    bl_idname = "ui_plus.opencv"
    bl_icon = "ops.generic.select_circle"
    bl_options = {'REGISTER'}

    bl_icon = "ops.generic.select_circle"

        
    def draw_settings(context, layout, tool):
        row = layout.row()
        op = row.operator("wm.opencv_operator", text="Capture", icon="OUTLINER_OB_CAMERA")

def register():
    bpy.utils.register_tool(OBJECT_MT_OpenCVPanel, separator=True, group=True)

def unregister():
    bpy.utils.unregister_tool(OBJECT_MT_OpenCVPanel)

if __name__ == "__main__":
    register()