change pose colors

libs/easy_dwpose/draw/openpose.py

@@ -36,28 +36,41 @@ def draw_bodypose(canvas, candidate, subset):
     ]
 
     colors = [
-        [255, 0, 0],
-        [255, 85, 0],
-        [255, 170, 0],
-        [255, 255, 0],
+        [0, 255, 255],
+        [255, 255, 255],
+
+        # arm-r
         [170, 255, 0],
-        [85, 255, 0],
+        [255, 255, 0],
+        [255, 170, 0],
+
+        # arm-l
+        [0, 85, 0],
         [0, 255, 0],
-        [0, 255, 85],
-        [0, 255, 170],
-        [0, 255, 255],
-        [0, 170, 255],
+        [0, 170, 0],
+
+        #leg-r
+        [255, 85, 0],
+        [255, 0, 0],
+        [255, 0, 85],
+
+        # leg-l
         [0, 85, 255],
         [0, 0, 255],
         [85, 0, 255],
+
         [170, 0, 255],
+        [0, 255, 170],
         [255, 0, 255],
-        [255, 0, 170],
-        [255, 0, 85],
+        [0, 170, 255],
     ]
 
+    
+
     for i in range(17):
         for n in range(len(subset)):
+            j = np.array(limbSeq[i]) - 1
+            #print(i, j)
             index = subset[n][np.array(limbSeq[i]) - 1]
             if -1 in index:
                 continue
@@ -68,7 +81,7 @@ def draw_bodypose(canvas, candidate, subset):
             length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
             angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
             polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
-            cv2.fillConvexPoly(canvas, polygon, colors[i])
+            cv2.fillConvexPoly(canvas, polygon, colors[j[1]])
 
     canvas = (canvas * 0.6).astype(np.uint8)
 

main.py

@@ -353,6 +353,7 @@ def prepare_inputs_inference(in_img, in_vid, frames, fps, dwpose, rembg_session,
             tpose.save("out/"+"tar_pose"+str(len(target_poses_cropped))+".png")
         target_poses_cropped.append(tpose)
         
+    target_poses_cropped[0].save("pose.png")
     return in_img, target_poses_cropped, in_pose, target_poses_coords, frames
 
 
@@ -1128,6 +1129,9 @@ def run_train_impl(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remov
     
     train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
 
+    gc.collect()
+    torch.cuda.empty_cache()
+
 def run_train(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True):
     run_train_impl(images, train_steps, modelId, bg_remove, resize_inputs)
 
@@ -1171,6 +1175,8 @@ def run_inference_impl(images, video_path, frames, train_steps=100, inference_st
     #results = [img_pad(img, img_width, img_height, True) for img in results]
     
     print("Done!")
+    gc.collect()
+    torch.cuda.empty_cache()
 
     return out_vid+'.webm', results, getThumbnails(results), target_poses_coords, getThumbnails(orig_frames)
 
@@ -1205,6 +1211,8 @@ def generate_frame(images, target_poses, train_steps=100, inference_steps=10, mo
     #results = [img_pad(img, img_width, img_height, True) for img in results]
     
     print("Done!")
+    gc.collect()
+    torch.cuda.empty_cache()
     
     results[0].save('result.png')