From 7fba9b84b0af95c0b83ee612b7c620058c86faf3 Mon Sep 17 00:00:00 2001
From: AEtheve <30652308+AEtheve@users.noreply.github.com>
Date: Tue, 2 May 2023 20:13:51 +0200
Subject: [PATCH] Added a mode to enable pose editing of a skeleton from the
 "Open Pose Editor" option

A new mode to the software that allows users to edit the pose of a skeleton from the "Open Pose Editor" option, using the image of the skeleton as input.
---
 electron_app/src/components/ControlNet.vue    |  32 ++-
 .../src/components_bare/PoseEditor.vue        | 190 ++++++++++++++++++
 2 files changed, 213 insertions(+), 9 deletions(-)
 create mode 100644 electron_app/src/components_bare/PoseEditor.vue

diff --git a/electron_app/src/components/ControlNet.vue b/electron_app/src/components/ControlNet.vue
index dac13b28..4dcc318b 100644
--- a/electron_app/src/components/ControlNet.vue
+++ b/electron_app/src/components/ControlNet.vue
@@ -6,7 +6,8 @@
             <!-- <p>Input Image:</p>
             <img class="gal_img" style="width: 90%;" src="https://colormadehappy.com/wp-content/uploads/2022/02/How-to-draw-a-cute-dog-6.jpg"/> 
          -->
-
+                <PoseEditor ref="pose_editor_canvas" style="height: calc(100% - 200px);" v-if="do_controlnet_preprocess=='No' && openpose" :canvas_width="600" :canvas_height="420" ></PoseEditor>
+            <span v-if="!(openpose && do_controlnet_preprocess=='No')">
             <div v-if="inp_img" @drop.prevent="onDragFile" @dragover.prevent class="image_area" :class="{ crosshair_cur  : is_inpaint }"  style="height: calc(100% - 200px);  border-radius: 16px; padding:5px;">
                 <ImageCanvas ref="inp_img_canvas" :is_inpaint="is_inpaint" :image_source="inp_img"  :is_disabled="!stable_diffusion.is_input_avail" canvas_id="ctrlnet" canvas_d_id="ctrlnetd" ></ImageCanvas>
             </div>
@@ -16,6 +17,8 @@
                 </center>
             </div>
 
+        </span>
+
             <div v-if="inp_img && stable_diffusion.is_input_avail" class="l_button" @click="open_input_image" >Change Image</div>
             <div v-if="inp_img && stable_diffusion.is_input_avail" class="l_button" @click="inp_img =''">Clear</div>
           
@@ -39,6 +42,7 @@
 
             <div v-if="stable_diffusion.is_input_avail" class="content_toolbox" style="margin-top:10px; margin-bottom:-10px;">
                 <div class="l_button button_medium button_colored" style="float:right ; " @click="generate_img2img" >Generate</div>
+                <div v-if="do_controlnet_preprocess=='No'" class="l_button button_medium button_colored" style="float:right ; " @click="open_pose_editor" >Open pose Editor</div>
                 <SDOptionsDropdown :options_model_values="this_object" :elements_hidden="['img_h' , 'img_w' ]" :elements_extra="['controlnet']"  > </SDOptionsDropdown>
             </div>
             <div v-else-if="stable_diffusion.generated_by=='controlnet'"  class="content_toolbox" style="margin-top:10px; margin-bottom:-10px;">
@@ -51,8 +55,6 @@
         </div>
 
         <div  class="right_half">
-
-        
             <div v-if="generated_images.length > 0 " >
 
                 <br> <br>
@@ -86,6 +88,7 @@
 <script>
 import ImageItem from '../components/ImageItem.vue'
 import ImageCanvas from '../components_bare/ImageCanvas.vue'
+import PoseEditor from '../components_bare/PoseEditor.vue'
 
 import LoaderModal from '../components_bare/LoaderModal.vue'
 import Vue from 'vue'
@@ -97,9 +100,8 @@ export default {
         app_state : Object   , 
         stable_diffusion : Object,
     },
-    components: { LoaderModal, ImageItem, ImageCanvas, SDOptionsDropdown },
+    components: { LoaderModal, ImageItem, ImageCanvas, PoseEditor, SDOptionsDropdown },
     mounted() {
-
     },
     computed:{
         is_sd_active(){
@@ -140,10 +142,12 @@ export default {
             selected_model : 'Default', 
             selected_control: "Body Pose",
             do_controlnet_preprocess: "Yes",
+            openpose: false,
+            stage_w :1000,
+            stage_h : 1000,
         };
     },
     methods: {
-
         generate_img2img(){
 
             let seed = 0;
@@ -156,17 +160,23 @@ export default {
                 Vue.$toast.default('You need to enter a prompt')
                 return;
             }
-                
 
-            if(!this.inp_img){
+
+            if(!this.inp_img && !this.openpose){
                 Vue.$toast.default('You need to add an input image')
                 return;
             }
                 
+            let input_image;
+            if (this.openpose && this.do_controlnet_preprocess=='No') {
+                input_image = window.ipcRenderer.sendSync('save_b64_image', this.$refs.pose_editor_canvas.get_img_b64());
+            } else {
+                input_image = window.ipcRenderer.sendSync('save_b64_image', this.$refs.inp_img_canvas.get_img_b64());
+            }
 
-            let input_image = window.ipcRenderer.sendSync('save_b64_image',  this.$refs.inp_img_canvas.get_img_b64()  );
             let input_image_with_mask;
             let mask_img;
+
             if(this.is_inpaint)
             { 
                 input_image_with_mask = window.ipcRenderer.sendSync('save_b64_image',  this.$refs.inp_img_canvas.get_img_mask_bg4() );
@@ -252,6 +262,10 @@ export default {
                 this.stable_diffusion.text_to_img(params, callbacks, 'controlnet');
         } , 
 
+        open_pose_editor(){
+            this.openpose = !this.openpose;
+        },
+
         open_input_image(){
             if( !this.stable_diffusion.is_input_avail)
                 return;
diff --git a/electron_app/src/components_bare/PoseEditor.vue b/electron_app/src/components_bare/PoseEditor.vue
new file mode 100644
index 00000000..63d5c750
--- /dev/null
+++ b/electron_app/src/components_bare/PoseEditor.vue
@@ -0,0 +1,190 @@
+<template>
+    <div style="height:100% ; width:100%; position:  relative;">
+        <div id="openpose_container">
+        <div id="openpose_canvas" ></div>
+    </div>
+   </div>
+</template>
+<script>
+import Konva from 'konva';
+
+function onVisible(element, callback) {
+  new IntersectionObserver((entries, observer) => {
+    entries.forEach(entry => {
+      if(entry.intersectionRatio > 0) {
+        callback(element);
+        observer;
+        // observer.disconnect();
+      }
+    });
+  }).observe(element);
+}
+
+export default {
+    name: 'PoseEditor',
+    props: {
+        canvas_width : Number,
+        canvas_height : Number,
+    },
+    components: {},
+    data() {
+        return {
+            stage_w : this.canvas_width,
+            stage_h : this.canvas_height,
+        }
+    },
+    mounted() {
+        this.init_state();
+        this.resize_stage();
+    },
+    computed : {
+    },
+    methods: {
+        init_state(){
+            this.stage_w = this.canvas_width;
+            this.stage_h = this.canvas_height;
+            
+            this.stage = new Konva.Stage({
+                container: 'openpose_canvas',
+                width: this.stage_w,
+                height: this.stage_h,
+            });
+
+            this.alpha_layer = new Konva.Layer({opacity	:1 });
+            
+            this.stage.add(this.alpha_layer);
+
+            let bg = new Konva.Rect({
+                                x: 0,
+                                y: 0,
+                                width: this.stage.width(),
+                                height: this.stage.height(),
+                                fill : "black",
+                            });
+            this.alpha_layer.add(bg);
+
+
+            function addVertex(alpha_layer, name, x, y, color){
+                let box = new Konva.Circle({
+                    x: x,
+                    y: y,
+                    radius: 10,
+                    fill: color,
+                    draggable: true,
+                    name: name
+                });
+                box.on('mouseover', function () {
+                    document.body.style.cursor = 'move';
+
+                });
+                box.on('mouseout', function () {
+                    document.body.style.cursor = 'default';
+                });
+                alpha_layer.add(box);
+                return box;
+            }
+
+            function addEdge(alpha_layer, box1, box2, color){
+                let line = new Konva.Line({
+                    points: [box1.x(), box1.y(), box2.x(), box2.y()],
+                    stroke: color,
+                    strokeWidth: 10,
+                    lineCap: 'round',
+                    lineJoin: 'round',
+                    opacity: 0.5,
+                });
+                alpha_layer.add(line);
+                line.setZIndex(1);
+
+                box1.on('dragmove', function () {
+                    line.points([box1.x(), box1.y(), box2.x(), box2.y()]);
+                });
+                box2.on('dragmove', function () {
+                    line.points([box1.x(), box1.y(), box2.x(), box2.y()]);
+                });
+            }
+
+            let vertex_join = [
+                ['left_shoulder', 'left_elbow'],
+                ['left_elbow', 'left_wrist'],
+                ['right_shoulder', 'right_elbow'],
+                ['right_elbow', 'right_wrist'],
+                ['left_shoulder', 'neck'],
+                ['right_shoulder', 'neck'],
+                ['neck', 'left_hip'],
+                ['neck', 'right_hip'],
+                ['left_hip', 'left_knee'],
+                ['left_knee', 'left_ankle'],
+                ['right_hip', 'right_knee'],
+                ['right_knee', 'right_ankle'],
+                ['nose', 'neck'],
+                ['nose', 'left_eye'],
+                ['nose', 'right_eye'],
+                ['left_eye', 'left_ear'],
+                ['right_eye', 'right_ear'],
+            ];
+
+            let point_positions = [
+                ['left_shoulder', 100, 100, 'rgb(255, 0, 0)'],
+                ['left_elbow', 90, 200, 'rgb(255, 0, 0)'],
+                ['left_wrist', 70, 300, 'rgb(255, 0, 0)'],
+                ['right_shoulder', 200, 100, 'rgb(0, 255, 0)'],
+                ['right_elbow', 210, 200, 'rgb(0, 255, 0)'],
+                ['right_wrist', 230, 300, 'rgb(0, 255, 0)'],
+                ['neck', 150, 100, 'rgb(0, 0, 255)'],
+                ['left_hip', 125, 200, 'rgb(0, 0, 255)'],
+                ['right_hip', 175, 200, 'rgb(0, 0, 255)'],
+                ['left_knee', 105, 300, 'rgb(0, 0, 255)'],
+                ['right_knee', 195, 300, 'rgb(0, 0, 255)'],
+                ['left_ankle', 90, 400, 'rgb(0, 0, 255)'],
+                ['right_ankle', 210, 400, 'rgb(0, 0, 255)'],
+                ['nose', 150, 50, 'rgb(255, 255, 0)'],
+                ['left_eye', 130, 30, 'rgb(255, 255, 0)'],
+                ['right_eye', 170, 30, 'rgb(255, 255, 0)'],
+                ['left_ear', 100, 50, 'rgb(255, 0, 255)'],
+                ['right_ear', 200, 50, 'rgb(255, 0, 255)'],
+            ];
+            
+            for (let i = 0; i < point_positions.length; i++) {
+                addVertex(this.alpha_layer, point_positions[i][0], point_positions[i][1], point_positions[i][2], point_positions[i][3]);
+            }
+            for (let i = 0; i < vertex_join.length; i++) {
+                let point1 = this.alpha_layer.find('.' + vertex_join[i][0])[0];
+                let point2 = this.alpha_layer.find('.' + vertex_join[i][1])[0];
+                let vertex_color = point_positions.find((element) => element[0] == vertex_join[i][0])[3];
+                addEdge(this.alpha_layer, point1, point2, vertex_color);
+            }
+
+            this.resize_stage()
+            
+
+            window.addEventListener('resize', this.resize_stage);
+            onVisible(document.querySelector("#openpose_container") , this.resize_stage )
+
+        },
+        resize_stage() {
+            let container = document.querySelector('#openpose_container');
+
+            // now we need to fit stage into parent container
+            let containerWidth = container.offsetWidth;
+            if(containerWidth < 10 )
+                return;
+
+            // but we also make the full scene visible
+            // so we need to scale all objects on canvas
+            let scale = containerWidth / this.stage_w;
+
+            this.stage.width(this.stage_w * scale - 5 );
+            this.stage.height(this.stage_h * scale);
+            this.stage.scale({ x: scale, y: scale });
+        },
+        get_img_b64(){
+            return this.stage.toDataURL();
+        },
+    },
+}
+</script>
+<style>
+</style>
+<style scoped>
+</style>
\ No newline at end of file