66import cv2
77import numpy as np
88import torch
9- from torch ._six import string_classes , int_classes
9+ from torch ._six import string_classes
1010
1111RED = (0 , 0 , 255 )
1212GREEN = (0 , 255 , 0 )
@@ -57,7 +57,7 @@ def collate_fn(batch):
5757 if elem .shape == (): # scalars
5858 py_type = float if elem .dtype .name .startswith ('float' ) else int
5959 return numpy_type_map [elem .dtype .name ](list (map (py_type , batch )))
60- elif isinstance (batch [0 ], int_classes ):
60+ elif isinstance (batch [0 ], int ):
6161 return torch .LongTensor (batch )
6262 elif isinstance (batch [0 ], float ):
6363 return torch .DoubleTensor (batch )
@@ -194,7 +194,7 @@ def vis_frame(frame, im_res, format='coco'):
194194 bg = img .copy ()
195195 cv2 .circle (bg , (int (cor_x / 2 ), int (cor_y / 2 )), 2 , p_color [n ], - 1 )
196196 # Now create a mask of logo and create its inverse mask also
197- transparency = max (0 , min (1 , kp_scores [n ]))
197+ transparency = max (0 , min (1 , kp_scores [n ]. item () ))
198198 img = cv2 .addWeighted (bg , transparency , img , 1 - transparency , 0 )
199199
200200 # Draw proposal score on the head
@@ -219,7 +219,7 @@ def vis_frame(frame, im_res, format='coco'):
219219 polygon = cv2 .ellipse2Poly ((int (mX ), int (mY )), (int (length / 2 ), int (stickwidth )), int (angle ), 0 , 360 , 1 )
220220 cv2 .fillConvexPoly (bg , polygon , line_color [i ])
221221 # cv2.line(bg, start_xy, end_xy, line_color[i], (2 * (kp_scores[start_p] + kp_scores[end_p])) + 1)
222- transparency = max (0 , min (1 , 0.5 * (kp_scores [start_p ] + kp_scores [end_p ])))
222+ transparency = max (0 , min (1 , 0.5 * (kp_scores [start_p ] + kp_scores [end_p ]). item () ))
223223 img = cv2 .addWeighted (bg , transparency , img , 1 - transparency , 0 )
224224 img = cv2 .resize (img , (width , height ), interpolation = cv2 .INTER_CUBIC )
225225 return img
0 commit comments