SHARE
TWEET

Untitled

a guest Dec 9th, 2019 104 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import io
  2. import cv2
  3. import numpy as np
  4. from google.cloud import firestore, storage
  5. import pickle
  6. from enum import Enum
  7. import math
  8. from modules.utils import get_intersection, perspective_M
  9. # db = firestore.Client()
  10. # client = storage.Client()
  11. #
  12. # in_bucket = client.get_bucket('recognition-frames')
  13. # blobs = in_bucket.list_blobs(prefix='rcg-client-00005/2019-10-16.15-54-19-00005')
  14. # blobs = in_bucket.list_blobs(prefix='rcg-client-00005/2019-11-30.09-34-57-00005')
  15. #
  16. # for blob in blobs:
  17. #         b = blob.download_as_string()
  18. #         nparr = np.fromstring(b, np.uint8)
  19. #         img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # cv2.IMREAD_COLOR in OpenCV 3.1
  20. #         print('Saving: ', blob.name)
  21. #         cv2.imwrite("images/big_session/" + blob.name.replace("/", ""), img_np)
  22. #
  23.  
  24. kpt_names = ['nose', 'neck',
  25.             'r_sho', 'r_elb', 'r_wri', 'l_sho', 'l_elb', 'l_wri',
  26.             'r_hip', 'r_knee', 'r_ank', 'l_hip', 'l_knee', 'l_ank',
  27.             'r_eye', 'l_eye',
  28.             'r_ear', 'l_ear']
  29.  
  30.  
  31. DIRECTION_FACTOR = 0.5
  32.  
  33. RIGHT_WRIST_INDEX = 4
  34. RIGHT_ELBOW_INDEX = 3
  35.  
  36. LEFT_WRIST_INDEX = 7
  37. LEFT_ELBOW_INDEX = 6
  38.  
  39. YELLOW = (255, 255, 0)
  40. TEAL = (0, 255, 255)
  41. RED = (0, 0, 255)
  42. GREEN = (36,255,12)
  43.  
  44. PROXIMITY_INFLUENCE = 1
  45. HORIZONTAL_DIST_INFLUENCE = 0.3
  46.  
  47. HALF_POINT_X = 1296 / 2
  48.  
  49. CART_FRONT_LEFT_CORNER = (490, 600)
  50. CART_FRONT_RIGHT_CORNER = (795, 595)
  51.  
  52. CART_BACK_LEFT_CORNER = (50, 720)
  53. CART_BACK_RIGHT_CORNER = (1250, 800)
  54.  
  55. left_cart_line = np.polyfit(CART_FRONT_LEFT_CORNER, CART_BACK_LEFT_CORNER, 1)
  56. right_cart_line = np.polyfit(CART_FRONT_RIGHT_CORNER, CART_BACK_RIGHT_CORNER, 1)
  57.  
  58. def keypoint_exists(keypoint):
  59.     return keypoint[0] > 0 and keypoint[1] > 0
  60.  
  61. def set_correct_bbox(coords, radius):
  62.     """ Finds the top left and bottom right coordinates of a
  63.        square based on it's center and radius
  64.  
  65.    Arguments:
  66.        coords {list} -- coordinates (x,y) of the center
  67.        radius {int} -- radius of the circle that would fit within the square
  68.  
  69.    Returns:
  70.        list -- (x1, y1, x2, y2), top left and bottom right coordinates of the square
  71.    """
  72.     x, y = coords
  73.     x1 = x - radius
  74.     x2 = x + radius
  75.     y1 = y - radius
  76.     y2 = y + radius
  77.     return int(x1), int(y1), int(x2), int(y2)
  78.  
  79. def draw_bbox_based_on_arm(img, wrist, elbow, min_distance, color):
  80.     vec_dir = wrist - elbow
  81.  
  82.     if min_distance <= 0:
  83.         prox_radius = 0
  84.     else:
  85.         prox_radius = (2 * int(700 / min_distance * 35 + 50)) * PROXIMITY_INFLUENCE
  86.  
  87.     horizontal_radius = vec_dir[0] * HORIZONTAL_DIST_INFLUENCE
  88.     radius = prox_radius + horizontal_radius
  89.  
  90.     bbox_center = wrist + DIRECTION_FACTOR * vec_dir
  91.     x1, y1, x2, y2 = set_correct_bbox(bbox_center, radius)
  92.  
  93.     start_point = (x1, y1)
  94.     end_point = (x2, y2)
  95.  
  96.     cv2.rectangle(img, start_point, end_point, color, 2)
  97.  
  98. if __name__ == "__main__":
  99.     filehandle = io.open("images/big_session/_out.pkl", 'rb')
  100.     # filehandle = io.open("images/out.pkl", 'rb')
  101.     unpickler = pickle.Unpickler(filehandle)
  102.     while filehandle.peek(1):
  103.         (filename, current_poses) = unpickler.load()
  104.         filename = filename.split('/', 1)
  105.         filename = filename[0] + '/' + "big_session/" + "rcg-client-00005" + filename[1].replace("/", "")
  106.         # filename = filename[0] + '/' + filename[1].replace("/", "")
  107.         img = cv2.imread(filename)
  108.  
  109.         split = filename.rsplit('.', 1)[0].rsplit('-', 2)[-2:]
  110.         try:
  111.             left = int(split[0])
  112.             right = int(split[1])
  113.         except AttributeError:
  114.             left = 0
  115.             right = 0
  116.         except ValueError:
  117.             left = 0
  118.             right = 0
  119.        
  120.         if left == 0:
  121.             min_distance = right
  122.         elif right == 0:
  123.             min_distance = left
  124.         else:
  125.             min_distance = min(left, right)
  126.  
  127.         for pose in current_poses:
  128.             if hasattr(pose, 'keypoints'):
  129.                 i = 0
  130.                 for keypt in pose.keypoints:  # Debug
  131.                     x, y = keypt
  132.                     if x > 0 and y > 0:
  133.                         cv2.circle(img, (keypt[0], keypt[1]), 4, (255, 0, 0), -1)
  134.                     i += 1
  135.                    
  136.                 cv2.circle(img, CART_FRONT_LEFT_CORNER, 10, YELLOW, 3) # Debug
  137.                 cv2.circle(img, CART_FRONT_RIGHT_CORNER, 10, YELLOW, 3) # Debug
  138.                 cv2.circle(img, CART_BACK_LEFT_CORNER, 10, YELLOW, 3) # Debug
  139.                 cv2.circle(img, CART_BACK_RIGHT_CORNER, 10, YELLOW, 3) # Debug
  140.  
  141.                 right_wrist, right_elbow = pose.keypoints[RIGHT_WRIST_INDEX], pose.keypoints[RIGHT_ELBOW_INDEX]
  142.                 if keypoint_exists(right_wrist) and keypoint_exists(right_elbow):
  143.                     cv2.circle(img, (right_wrist[0], right_wrist[1]), 10, YELLOW, 3) # Debug
  144.                     cv2.circle(img, (right_elbow[0], right_elbow[1]), 10, YELLOW, 3) # Debug
  145.                     draw_bbox_based_on_arm(img, right_wrist, right_elbow, min_distance, YELLOW)
  146.  
  147.                 elif keypoint_exists(right_wrist):
  148.                     cv2.circle(img, (right_wrist[0], right_wrist[1]), 10, RED, 3) # Debug
  149.                     draw_bbox_based_on_arm(img, right_wrist, (0, 0), min_distance, RED)
  150.    
  151.                 left_wrist, left_elbow = pose.keypoints[LEFT_WRIST_INDEX], pose.keypoints[LEFT_ELBOW_INDEX]
  152.                 if keypoint_exists(left_wrist) and keypoint_exists(left_elbow):
  153.                     cv2.circle(img, (left_wrist[0], left_wrist[1]), 10, TEAL, 3) # Debug
  154.                     cv2.circle(img, (left_elbow[0], left_elbow[1]), 10, TEAL, 3) # Debug
  155.                     draw_bbox_based_on_arm(img, left_wrist, left_elbow, min_distance, TEAL)
  156.                 elif keypoint_exists(left_wrist):
  157.                     cv2.circle(img, (left_wrist[0], left_wrist[1]), 10, GREEN, 3) # Debug
  158.                     draw_bbox_based_on_arm(img, left_wrist, (0, 0), min_distance, GREEN)
  159.  
  160.                 # Fall back when no pose detected
  161.                 # if not (drew_left or drew_right) and min_distance != 0:
  162.                 #     intersection, _ = get_intersection(left, right)
  163.                 #     if intersection:
  164.                 #         intersection = np.array([[intersection]], dtype='float32')
  165.                 #         intersection = cv2.perspectiveTransform(intersection, perspective_M)
  166.                 #         p = tuple(intersection[0][0])
  167.                 #         abs_center = [p[0], p[1]]
  168.                 #         radius = (2 * int(700 / min_distance * 35 + 50))                        
  169.                 #         x1, y1, x2, y2 = set_correct_bbox(abs_center, radius)
  170.                 #         cv2.rectangle(img, (x1, y1), (x2, y2), RED, 2)
  171.  
  172.         cv2.imshow("bla", img)
  173.         cv2.waitKey(0)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top