SHARE
TWEET

Untitled

a guest Dec 9th, 2019 78 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 = (300, 10)
  53. CART_BACK_RIGHT_CORNER = (900, 10)
  54.  
  55. def keypoint_exists(keypoint):
  56.     return keypoint[0] > 0 and keypoint[1] > 0
  57.  
  58. def set_correct_bbox(coords, radius):
  59.     """ Finds the top left and bottom right coordinates of a
  60.        square based on it's center and radius
  61.  
  62.    Arguments:
  63.        coords {list} -- coordinates (x,y) of the center
  64.        radius {int} -- radius of the circle that would fit within the square
  65.  
  66.    Returns:
  67.        list -- (x1, y1, x2, y2), top left and bottom right coordinates of the square
  68.    """
  69.     x, y = coords
  70.     x1 = x - radius
  71.     x2 = x + radius
  72.     y1 = y - radius
  73.     y2 = y + radius
  74.     return int(x1), int(y1), int(x2), int(y2)
  75.  
  76. def get_forearm_active_probability(img, wrist, elbow):
  77.     vec_dir = wrist - elbow
  78.     angle = abs(math.degrees(math.atan(float(vec_dir[1] / vec_dir[0]))))
  79.  
  80.     cv2.circle(img, CART_FRONT_LEFT_CORNER, 10, YELLOW, 3) # Debug
  81.     cv2.circle(img, CART_FRONT_RIGHT_CORNER, 10, YELLOW, 3) # Debug
  82.  
  83.     cv2.circle(img, CART_BACK_LEFT_CORNER, 10, YELLOW, 3)
  84.     cv2.circle(img, CART_BACK_RIGHT_CORNER, 10, YELLOW, 3)
  85.  
  86.     cv2.putText(img, str(angle),
  87.         (wrist[0] + 10, wrist[1] + 10),
  88.         cv2.FONT_HERSHEY_SIMPLEX,
  89.         0.5,
  90.         (255, 255, 255),
  91.         2)
  92.  
  93.     return img
  94.  
  95. def draw_bbox_based_on_arm(img, wrist, elbow, min_distance, color):
  96.     vec_dir = wrist - elbow
  97.  
  98.     if min_distance <= 0:
  99.         prox_radius = 0
  100.     else:
  101.         prox_radius = (2 * int(700 / min_distance * 35 + 50)) * PROXIMITY_INFLUENCE
  102.  
  103.     horizontal_radius = vec_dir[0] * HORIZONTAL_DIST_INFLUENCE
  104.     radius = prox_radius + horizontal_radius
  105.  
  106.     bbox_center = wrist + DIRECTION_FACTOR * vec_dir
  107.     x1, y1, x2, y2 = set_correct_bbox(bbox_center, radius)
  108.  
  109.     start_point = (x1, y1)
  110.     end_point = (x2, y2)
  111.  
  112.     img = get_forearm_active_probability(img, wrist, elbow)
  113.  
  114.     cv2.rectangle(img, start_point, end_point, color, 2)
  115.  
  116. if __name__ == "__main__":
  117.     filehandle = io.open("images/big_session/_out.pkl", 'rb')
  118.     # filehandle = io.open("images/out.pkl", 'rb')
  119.     unpickler = pickle.Unpickler(filehandle)
  120.     while filehandle.peek(1):
  121.         (filename, current_poses) = unpickler.load()
  122.         filename = filename.split('/', 1)
  123.         filename = filename[0] + '/' + "big_session/" + "rcg-client-00005" + filename[1].replace("/", "")
  124.         # filename = filename[0] + '/' + filename[1].replace("/", "")
  125.         img = cv2.imread(filename)
  126.  
  127.         split = filename.rsplit('.', 1)[0].rsplit('-', 2)[-2:]
  128.         try:
  129.             left = int(split[0])
  130.             right = int(split[1])
  131.         except AttributeError:
  132.             left = 0
  133.             right = 0
  134.         except ValueError:
  135.             left = 0
  136.             right = 0
  137.        
  138.         if left == 0:
  139.             min_distance = right
  140.         elif right == 0:
  141.             min_distance = left
  142.         else:
  143.             min_distance = min(left, right)
  144.  
  145.         for pose in current_poses:
  146.             if hasattr(pose, 'keypoints'):
  147.                 i = 0
  148.                 for keypt in pose.keypoints:  # Debug
  149.                     x, y = keypt
  150.                     if x > 0 and y > 0:
  151.                         cv2.circle(img, (keypt[0], keypt[1]), 4, (255, 0, 0), -1)
  152.                     i += 1
  153.  
  154.                 right_wrist, right_elbow = pose.keypoints[RIGHT_WRIST_INDEX], pose.keypoints[RIGHT_ELBOW_INDEX]
  155.                 if keypoint_exists(right_wrist) and keypoint_exists(right_elbow):
  156.                     cv2.circle(img, (right_wrist[0], right_wrist[1]), 10, YELLOW, 3) # Debug
  157.                     cv2.circle(img, (right_elbow[0], right_elbow[1]), 10, YELLOW, 3) # Debug
  158.                     draw_bbox_based_on_arm(img, right_wrist, right_elbow, min_distance, YELLOW)
  159.                     # TODO Detect if relevant
  160.  
  161.                 elif keypoint_exists(right_wrist):
  162.                     cv2.circle(img, (right_wrist[0], right_wrist[1]), 10, RED, 3) # Debug
  163.                     draw_bbox_based_on_arm(img, right_wrist, (0, 0), min_distance, RED)
  164.    
  165.                 left_wrist, left_elbow = pose.keypoints[LEFT_WRIST_INDEX], pose.keypoints[LEFT_ELBOW_INDEX]
  166.                 if keypoint_exists(left_wrist) and keypoint_exists(left_elbow):
  167.                     cv2.circle(img, (left_wrist[0], left_wrist[1]), 10, TEAL, 3) # Debug
  168.                     cv2.circle(img, (left_elbow[0], left_elbow[1]), 10, TEAL, 3) # Debug
  169.                     draw_bbox_based_on_arm(img, left_wrist, left_elbow, min_distance, TEAL)
  170.                 elif keypoint_exists(left_wrist):
  171.                     cv2.circle(img, (left_wrist[0], left_wrist[1]), 10, GREEN, 3) # Debug
  172.                     draw_bbox_based_on_arm(img, left_wrist, (0, 0), min_distance, GREEN)
  173.                     # TODO Detect if relevant
  174.  
  175.                 # Fall back when no pose detected
  176.                 # if not (drew_left or drew_right) and min_distance != 0:
  177.                 #     intersection, _ = get_intersection(left, right)
  178.                 #     if intersection:
  179.                 #         intersection = np.array([[intersection]], dtype='float32')
  180.                 #         intersection = cv2.perspectiveTransform(intersection, perspective_M)
  181.                 #         p = tuple(intersection[0][0])
  182.                 #         abs_center = [p[0], p[1]]
  183.                 #         radius = (2 * int(700 / min_distance * 35 + 50))                        
  184.                 #         x1, y1, x2, y2 = set_correct_bbox(abs_center, radius)
  185.                 #         cv2.rectangle(img, (x1, y1), (x2, y2), RED, 2)
  186.  
  187.         cv2.imshow("bla", img)
  188.         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
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!
 
Top