Update util.py

util.py

@@ -306,65 +306,57 @@ def handDetect(candidate, subset, oriImg):
     '''
     return detect_result
 
-def drawStickmodel(oriImg,x_ytupple,x_y_sticks,export_edges,export_peaks):
+def drawStickmodel(oriImg, x_ytupple, x_y_sticks, export_edges, export_peaks):
     canvas = copy.deepcopy(oriImg)
 
+    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0],
+              [85, 255, 0], [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255],
+              [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], [170, 0, 255],
+              [255, 0, 255], [255, 0, 170], [255, 0, 85], [255,255,0], [255,255,85],
+              [255,255,170], [255,255,255],[170,255,255],[85,255,255],[0,255,255]]
+    stickwidth = 4
 
-    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], 
-            [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], 
-            [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85], [255,255,0], [255,255,85], [255,255,170],
-                [255,255,255],[170,255,255],[85,255,255],[0,255,255]]
-    stickwidth=4
-
-    for idx,(mX,mY,angle,length) in enumerate(x_y_sticks):
+    for idx, (mX, mY, angle, length) in enumerate(x_y_sticks):
         cur_canvas = canvas.copy()
-        # print(f'new cv2.ellipse2Poly((int({mY}), int({mX})), (int({length} / 2), {stickwidth}), int({angle}), 0, 360, 1)')
-        polygon = cv2.ellipse2Poly((int(mX), int(mY)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+        polygon = cv2.ellipse2Poly((int(mX), int(mY)), (int(length / 2), stickwidth),
+                                   int(angle), 0, 360, 1)
         cv2.fillConvexPoly(cur_canvas, polygon, colors[idx])
         canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
 
-
-
-    for idx,(x,y) in enumerate(x_ytupple):
+    for idx, (x, y) in enumerate(x_ytupple):
         cv2.circle(canvas, (int(x), int(y)), 4, colors[idx], thickness=-1)
-        
 
-    ## Handpose
+    # Handpose
     fig = Figure(figsize=plt.figaspect(canvas))
     fig.subplots_adjust(0, 0, 1, 1)
-    fig.subplots_adjust(bottom=0, top=1, left=0, right=1)
-    bg = FigureCanvas(fig)
     ax = fig.subplots()
     ax.axis('off')
     ax.imshow(canvas)
 
-    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
-                [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
+    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9],
+             [9, 10], [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16],
+             [0, 17], [17, 18], [18, 19], [19, 20]]
 
     for both_hand_edges in export_edges:
-        for (ie,(x1, y1),(x2, y2)) in both_hand_edges:
-            # print(f'new ax.plot([{x1}, {x2}], [{y1}, {y2}], color=matplotlib.colors.hsv_to_rgb([ie/float({len(edges)}), 1.0, 1.0]))')
-            ax.plot([x1, x2], [y1, y2], color=matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0]))
-
-    width, height = ax.figure.get_size_inches() * ax.figure.get_dpi()
+        for (ie, (x1, y1), (x2, y2)) in both_hand_edges:
+            ax.plot([x1, x2], [y1, y2],
+                    color=matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0]))
 
     for both_hand_peaks in export_peaks:
-        for (x,y,text) in both_hand_peaks:
-            # print(f"new ax.plot({x}, {y}, 'r.')")
+        for (x, y, text) in both_hand_peaks:
             ax.plot(x, y, 'r.')
 
-
-
-    # print(f'NEW width = {width}, height={height}')
+    # Convert figure to numpy array
+    bg = FigureCanvas(fig)
     bg.draw()
 
-    canvas = np.fromstring(bg.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
-
-    ####
-
+    width, height = fig.get_size_inches() * fig.get_dpi()
+    buf = bg.buffer_rgba()
+    canvas = np.frombuffer(buf, dtype=np.uint8).reshape(int(height), int(width), 4)
+    canvas = canvas[:, :, :3]  # keep only RGB
 
-    # cv2.imwrite('C:/Users/spsar/Downloads/MVI_5177.MOV-transformed/MVI_5177.MOV-GaussianBlur/MVI_5177.MOV-14-modified.jpg', canvas) 
-    return cv2.resize(canvas,(math.ceil(width),math.ceil(height)))
+    plt.close(fig)  # clean up
+    return cv2.resize(canvas, (math.ceil(width), math.ceil(height)))
 
 def draw_bar_plot_below_image(image, predictions, title, origImg):
   """