Upload 10 files

microprograms/errors/angles_micro_programs.py

@@ -0,0 +1,107 @@
+import torch
+import numpy as np
+import math
+import cv2
+import sys, os
+from matplotlib import image
+from matplotlib import pyplot as plt
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+from math import atan
+
+# KEYPOINT_INDEXES = {
+#     0: 'r ankle',
+#     1: 'r knee',
+#     2: 'r hip',
+#     3: 'l hip',
+#     4: 'l knee',
+#     5: 'l ankle',
+#     6: 'pelvis',
+#     7: 'thorax',
+#     8: 'upper neck',
+#     9: 'head',
+#     10: 'r wrist',
+#     11: 'r elbow',
+#     12: 'r shoulder',
+#     13: 'l shoulder',
+#     14: 'l elbow',
+#     15: 'l wrist',
+# }
+
+def slope(x1, y1, x2, y2):
+    if x1 == x2:
+        return "undefined"
+    return (y2-y1)/(x2-x1)
+
+# Function to find the
+# angle between two lines
+def findAngle(M1, M2):
+    vertical_line = False
+    if M1 == "undefined":
+        M1 = 0
+        vertical_line = True
+    if M2 == "undefined":
+        M2 = 0
+        vertical_line = True
+    PI = 3.14159265
+     
+    # Store the tan value  of the angle
+    angle = abs((M2 - M1) / (1 + M1 * M2))
+ 
+    # Calculate tan inverse of the angle
+    ret = atan(angle)
+ 
+    # Convert the angle from
+    # radian to degree
+    val = (ret * 180) / PI
+ 
+    # Print the result
+    #print(round(val, 4))
+    if vertical_line:
+        return 90 - round(val,4)
+    return (round(val, 4))
+
+# np.array(pose_pred)[0][i]
+def applyFeetApartError(filepath, pose_pred=None, diver_detector=None, pose_model=None):
+    if pose_pred is None and filepath != "":
+        diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model)
+    if pose_pred is not None:
+        pose_pred = np.array(pose_pred)[0]
+        average_knee = [np.mean((pose_pred[4][0], pose_pred[1][0])), np.mean((pose_pred[4][1], pose_pred[1][1]))]
+        vector1 = [pose_pred[5][0] - average_knee[0], pose_pred[5][1] - average_knee[1]]
+        vector2 = [pose_pred[0][0] - average_knee[0], pose_pred[0][1] - average_knee[1]]
+        unit_vector_1 = vector1 / np.linalg.norm(vector1)
+        unit_vector_2 = vector2 / np.linalg.norm(vector2)
+        dot_product = np.dot(unit_vector_1, unit_vector_2)
+        angle = math.degrees(np.arccos(dot_product))
+        # left_leg_slope = slope(np.array(pose_pred)[0][4][0], np.array(pose_pred)[0][4][1], np.array(pose_pred)[0][5][0], np.array(pose_pred)[0][5][1])
+        # right_leg_slope = slope(np.array(pose_pred)[0][4][0], np.array(pose_pred)[0][4][1], np.array(pose_pred)[0][0][0], np.array(pose_pred)[0][0][1])
+        # angle = findAngle(left_leg_slope, right_leg_slope)
+        return angle
+    else:
+        # print('pose_pred is None')
+        return None
+
+#position: somersault or twist
+def applyPositionTightnessError(filepath, pose_pred=None, diver_detector=None, pose_model=None):
+    if pose_pred is None and filepath != "":
+        diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model)
+    if pose_pred is not None:
+        pose_pred = np.array(pose_pred)[0]
+        vector1 = [pose_pred[7][0] - pose_pred[2][0], pose_pred[7][1] - pose_pred[2][1]]
+        
+        vector2 = [pose_pred[1][0] - pose_pred[2][0], pose_pred[1][1] - pose_pred[2][1]]
+        unit_vector_1 = vector1 / np.linalg.norm(vector1)
+        unit_vector_2 = vector2 / np.linalg.norm(vector2)
+        dot_product = np.dot(unit_vector_1, unit_vector_2)
+        angle = math.degrees(np.arccos(dot_product))
+        # upper_body = slope(np.array(pose_pred)[0][2][0], np.array(pose_pred)[0][2][1], np.array(pose_pred)[0][7][0], np.array(pose_pred)[0][7][1])
+        # lower_body = slope(np.array(pose_pred)[0][2][0], np.array(pose_pred)[0][2][1], np.array(pose_pred)[0][1][0], np.array(pose_pred)[0][1][1])
+        # angle = findAngle(upper_body, lower_body)
+    # if position == 1:
+    #     angle = 180 - angle
+        return angle
+    else:
+        # print('pose_pred is None')
+        return None
+
+

microprograms/errors/distance_from_springboard_micro_program.py

@@ -0,0 +1,290 @@
+# import sys, os, distutils.core
+
+# # os.system('python -m pip install pyyaml==5.3.1')
+# # dist = distutils.core.run_setup("./detectron2/setup.py")
+# # temp = ' '.join([f"'{x}'" for x in dist.install_requires])
+# # cmd = "python -m pip install {0}".format(temp)
+# # os.system(cmd)
+# sys.path.insert(0, os.path.abspath('./detectron2'))
+
+# import detectron2
+# import cv2
+
+# from detectron2.utils.logger import setup_logger
+# setup_logger()
+
+# # from detectron2.modeling import build_model
+# from detectron2 import model_zoo
+# from detectron2.engine import DefaultPredictor
+# from detectron2.config import get_cfg
+# from detectron2.utils.visualizer import Visualizer
+# from detectron2.data import MetadataCatalog, DatasetCatalog
+# from detectron2.utils.visualizer import Visualizer
+# from detectron2.checkpoint import DetectionCheckpointer
+# from detectron2.data.datasets import register_coco_instances
+
+# cfg = get_cfg()
+# cfg.OUTPUT_DIR = "./output/springboard/"
+# # model = build_model(cfg)  # returns a torch.nn.Module
+# cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
+# cfg.DATASETS.TEST = ()
+# cfg.DATALOADER.NUM_WORKERS = 2
+# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
+# cfg.SOLVER.IMS_PER_BATCH = 2  # This is the real "batch size" commonly known to deep learning people
+# cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
+# cfg.SOLVER.MAX_ITER = 300    # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
+# cfg.SOLVER.STEPS = []        # do not decay learning rate
+# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # The "RoIHead batch size". 128 is faster, and good enough for this toy dataset (default: 512)
+# cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+# cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # path to the model we just trained
+# cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a custom testing threshold
+# predictor = DefaultPredictor(cfg)
+# register_coco_instances("springboard_trains", {}, "./coco_annotations/springboard/train.json", "../data/Boards/spring")
+# register_coco_instances("springboard_vals", {}, "./coco_annotations/springboard/val.json", "../data/Boards/spring")
+
+# from detectron2.utils.visualizer import ColorMode
+# splash_metadata = MetadataCatalog.get('springboard_vals')
+# dataset_dicts = DatasetCatalog.get("springboard_vals")
+
+# outputs_array = []
+# for d in dataset_dicts:
+#     im = cv2.imread(d["file_name"])
+#     outputs = predictor(im)
+#     outputs_array.append(outputs)  # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
+#     v = Visualizer(im[:, :, ::-1],
+#                    metadata=splash_metadata,
+#                    scale=0.5,
+#                    instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models
+#     )
+#     out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+#     img = out.get_image()[:, :, ::-1]
+#     filename = os.path.join("./output/", d["file_name"][3:])
+#     print(filename)
+#     if not cv2.imwrite(filename, img):
+#       print('no image written')
+
+import torch
+import numpy as np
+import math
+import cv2
+import sys, os
+from matplotlib import image
+from matplotlib import pyplot as plt
+from models.detectron2.springboard_detector_setup import get_springboard_detector
+from models.detectron2.platform_detector_setup import get_platform_detector
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+
+# springboard MICRO PROGRAM
+
+# returns "left" or "right" depending on whether the board is on the left or right side of the frame
+def find_which_side_board_on(output):
+    pred_classes = output['instances'].pred_classes.cpu().numpy()
+    platforms = np.where(pred_classes == 0)[0]
+    scores = output['instances'].scores[platforms]
+    if len(scores) == 0:
+      return
+    pred_masks = output['instances'].pred_masks[platforms]
+    max_instance = torch.argmax(scores)
+    pred_mask = np.array(pred_masks[max_instance].cpu()) 
+    for i in range(len(pred_mask[0])//2):
+        if sum(pred_mask[:, i]) != 0:
+            return "left"
+        elif sum(pred_mask[:, len(pred_mask[0]) - i - 1]) != 0:
+            return "right"
+    return None
+
+def board_end(output, board_side=None):
+    # pred_classes = output['instances'].pred_classes.cpu().numpy()
+    # splashes = np.where(pred_classes == 0)[0]
+    # scores = output['instances'].scores[splashes]
+    # if len(scores) == 0:
+    #     return
+    # pred_masks = output['instances'].pred_masks[splashes]
+    # max_instance = torch.argmax(scores)
+    # pred_mask = pred_masks[max_instance] # splash instance with highest confidence
+    pred_classes = output['instances'].pred_classes.cpu().numpy()
+    platforms = np.where(pred_classes == 0)[0]
+    scores = output['instances'].scores[platforms]
+    if len(scores) == 0:
+      return
+    pred_masks = output['instances'].pred_masks[platforms]
+    max_instance = torch.argmax(scores)
+    pred_mask = np.array(pred_masks[max_instance].cpu()) # splash instance with highest confidence
+    # need to figure out whether springboard is on left or right side of frame, then need to find where the edge is
+    if board_side is None:
+        board_side = find_which_side_board_on(output)
+    if board_side == "left":
+      for i in range(len(pred_mask[0]) - 1, -1, -1):
+        if sum(pred_mask[:, i]) != 0:
+          trues = np.where(pred_mask[:, i])[0]
+          return (i, min(trues))
+    if board_side == "right":
+      for i in range(len(pred_mask[0])):
+        if sum(pred_mask[:, i]) != 0:
+          trues = np.where(pred_mask[:, i])[0]
+          return (i, min(trues))
+    return None
+
+def draw_board_end_coord(im, coord):
+    print("hello, im in the drawing func")
+    image = cv2.circle(im, (int(coord[0]),int(coord[1])), radius=10, color=(0, 0, 255), thickness=-1)
+    filename = os.path.join("./output/board_end/", d["file_name"][3:])
+    print(filename)
+    if not cv2.imwrite(filename, image):
+        print(filename)
+        print("file failed to write")
+
+# loops over each image, plots a point for the end of board
+# i = 0
+# for d in dataset_dicts:
+#   im = cv2.imread(d["file_name"])
+#   outputs = predictor(im)
+
+#   # to draw a point on co-ordinate (200,300)
+#   coord = board_end(outputs)
+#   if coord == None:
+#     continue
+#   # plt.plot(coord[0], coord[1], marker='v', color="white")
+#   draw_board_end_coord(im, coord)
+#   i+=1
+
+## TODO: ADD POSE ESTIMATOR, AND CALCULATE DISTANCE FROM BOARD
+# PLOT RESULTS OF ONE FULL DIVE 
+# KEYPOINT_INDEXES = {
+#     0: 'r ankle',
+#     1: 'r knee',
+#     2: 'r hip',
+#     3: 'l hip',
+#     4: 'l knee',
+#     5: 'l ankle',
+#     6: 'pelvis',
+#     7: 'thorax',
+#     8: 'upper neck',
+#     9: 'head',
+#     10: 'r wrist',
+#     11: 'r elbow',
+#     12: 'r shoulder',
+#     13: 'l shoulder',
+#     14: 'l elbow',
+#     15: 'l wrist',
+# }
+
+# demo_image = '../data/Boards/spring/img_17_10_00014517.jpg'
+# im = cv2.imread(demo_image)
+
+# pose_pred = get_pose_estimation(demo_image)
+# print("pose_pred", pose_pred)
+
+# predictor = get_springboard_detector()
+# outputs = predictor(im)
+# # to draw a point on co-ordinate (200,300)
+# coord = board_end(outputs)
+
+def draw_two_coord(im, coord1, coord2, filename):
+    print("hello, im in the drawing func")
+    image = cv2.circle(im, (int(coord1[0]),int(coord1[1])), radius=5, color=(0, 0, 255), thickness=-1)
+    image = cv2.circle(image, (int(coord2[0]),int(coord2[1])), radius=5, color=(0, 255, 0), thickness=-1)
+    print(filename)
+    if not cv2.imwrite(filename, image):
+        print(filename)
+        print("file failed to write")
+
+# draw_two_coord(im, coord, np.array(pose_pred)[0][5], filename==os.path.join("./output/board_end/", demo_image[3:]))
+
+# print("pose_pred.shape", np.array(pose_pred).shape)
+# print("coord.shape", np.array(coord).shape)
+# print("DISTANCE BETWEEN END BOARD AND LEFT ANKLE:", math.dist(np.array(pose_pred)[0][5], np.array(coord)))
+
+def calculate_distance_from_springboard_for_one_frame(filepath, visualize=False, dive_folder_num="", springboard_detector=None, pose_pred=None, pose_model=None, board_end_coord=None, board_side=None):
+    if springboard_detector is None:
+        springboard_detector = get_springboard_detector()
+    if pose_pred is None:
+        diver_box, pose_pred = get_pose_estimation(filepath, pose_model=pose_model)
+    im = cv2.imread(filepath)
+    outputs = springboard_detector(im)
+    if board_end_coord is None:
+        board_end_coord = board_end(outputs, board_side=board_side)
+    minDist = None
+    if board_end_coord is not None and pose_pred is not None and len(board_end_coord) == 2:
+        minDist = float('inf')
+        for i in range(len(np.array(pose_pred)[0])):
+            distance = math.dist(np.array(pose_pred)[0][i], np.array(board_end_coord))
+            if distance < minDist:
+                minDist = distance
+                minJoint = i
+        if visualize:
+            file_name = filepath.split('/')[-1]
+            folder = "./output/data/distance_from_board/{}".format(dive_folder_num)
+            out_filename = os.path.join(folder, file_name)
+            if not os.path.exists(folder):
+                os.makedirs(folder)
+            draw_two_coord(im, board_end_coord, np.array(pose_pred)[0][minJoint], filename=out_filename)
+    ## more verbose
+    # else:
+        # print("springboard or diver not detected in", filepath)
+        # if board_end_coord is None:
+        #     print("springboard not detected in", filepath)
+        # if pose_pred is None:
+        #     print("diver not detected in", filepath)
+    return minDist
+
+def calculate_distance_from_platform_for_one_frame(filepath, im=None, visualize=False, dive_folder_num="", platform_detector=None, pose_pred=None, diver_detector=None, pose_model=None, board_end_coord=None, board_side=None):
+    if platform_detector is None:
+        platform_detector = get_platform_detector()
+    if pose_pred is None:
+        diver_box, pose_pred = get_pose_estimation(filepath, image_bgr=im, diver_detector=diver_detector, pose_model=pose_model)
+    if im is None and filepath != "":
+        im = cv2.imread(filepath)
+    if board_end_coord is None:
+        outputs = platform_detector(im)
+        board_end_coord = board_end(outputs, board_side=board_side)
+    minDist = None
+    if board_end_coord is not None and pose_pred is not None and len(board_end_coord) == 2:
+        minDist = float('inf')
+        for i in range(len(np.array(pose_pred)[0])):
+            distance = math.dist(np.array(pose_pred)[0][i], np.array(board_end_coord))
+            if distance < minDist:
+                minDist = distance
+                minJoint = i
+        if visualize:
+            file_name = filepath.split('/')[-1]
+            folder = "./output/data/distance_from_board/{}".format(dive_folder_num)
+            out_filename = os.path.join(folder, file_name)
+            if not os.path.exists(folder):
+                os.makedirs(folder)
+            draw_two_coord(im, board_end_coord, np.array(pose_pred)[0][minJoint], filename=out_filename)
+    ## more verbose
+    # else:
+        # print("platform or diver not detected in", filepath)
+        # if board_end_coord is None:
+        #     print("platform not detected in", filepath)
+        # if pose_pred is None:
+        #     print("diver not detected in", filepath)
+    return minDist
+
+# distances = []
+# directory = "./FineDiving/datasets/FINADiving_MTL_256s/17/73/"
+# file_names = os.listdir(directory)
+# for file_name in file_names:
+#     path = os.path.join(directory, file_name)
+#     pose_pred = get_pose_estimation(path)
+#     print("PATH IM_17_73:", path)
+#     im = cv2.imread(path)
+#     outputs = predictor(im)
+#     coord = board_end(outputs)
+#     if coord is not None and pose_pred is not None and len(coord) == 2:
+#       distance = math.dist(np.array(pose_pred)[0][5], np.array(coord))
+#       if distance is None:
+#         distances.append(0)
+#       else:
+#         distances.append(distance)
+#         filename = os.path.join("./output/data/img_17_73/", file_name)
+#         draw_two_coord(im, coord, np.array(pose_pred)[0][5], filename=filename)
+#     else:
+#       distances.append(0)
+
+# plt.plot(range(len(distances)), distances)
+# plt.savefig('./output/data/img_17_73/img_17_73_board_dist_graph.png')
+
+
+

microprograms/errors/feet_apart.py

@@ -0,0 +1,20 @@
+import numpy as np
+import math
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+
+def applyFeetApartError(filepath, pose_pred=None, diver_detector=None, pose_model=None):
+    if pose_pred is None and filepath != "":
+        diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model)
+    if pose_pred is not None:
+        pose_pred = np.array(pose_pred)[0]
+        average_knee = [np.mean((pose_pred[4][0], pose_pred[1][0])), np.mean((pose_pred[4][1], pose_pred[1][1]))]
+        vector1 = [pose_pred[5][0] - average_knee[0], pose_pred[5][1] - average_knee[1]]
+        vector2 = [pose_pred[0][0] - average_knee[0], pose_pred[0][1] - average_knee[1]]
+        unit_vector_1 = vector1 / np.linalg.norm(vector1)
+        unit_vector_2 = vector2 / np.linalg.norm(vector2)
+        dot_product = np.dot(unit_vector_1, unit_vector_2)
+        angle = math.degrees(np.arccos(dot_product))
+        return angle
+    else:
+        # print('pose_pred is None')
+        return None

microprograms/errors/over_rotation.py

@@ -0,0 +1,19 @@
+import numpy as np
+import math
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+
+def over_rotation(filepath, pose_pred=None, diver_detector=None, pose_model=None):
+    if pose_pred is None and filepath != "":
+        diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model)
+    if pose_pred is not None:
+        pose_pred = np.array(pose_pred)[0]
+        vector1 = [(pose_pred[0][0] - pose_pred[2][0]), 0-(pose_pred[0][1] - pose_pred[2][1])]        
+        vector2 = [-1, 0]
+        unit_vector_1 = vector1 / np.linalg.norm(vector1)
+        unit_vector_2 = vector2 / np.linalg.norm(vector2)
+        dot_product = np.dot(unit_vector_1, unit_vector_2)
+        angle = math.degrees(np.arccos(dot_product))
+        return angle
+    else:
+        # print('pose_pred is None')
+        return None

microprograms/errors/position_tightness.py

@@ -0,0 +1,20 @@
+import numpy as np
+import math
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+
+def applyPositionTightnessError(filepath, pose_pred=None):
+    if pose_pred is None and filepath != "":
+        diver_box, pose_pred = get_pose_estimation(filepath)
+    if pose_pred is not None:
+        pose_pred = np.array(pose_pred)[0]
+        vector1 = [pose_pred[7][0] - pose_pred[2][0], pose_pred[7][1] - pose_pred[2][1]]
+        
+        vector2 = [pose_pred[1][0] - pose_pred[2][0], pose_pred[1][1] - pose_pred[2][1]]
+        unit_vector_1 = vector1 / np.linalg.norm(vector1)
+        unit_vector_2 = vector2 / np.linalg.norm(vector2)
+        dot_product = np.dot(unit_vector_1, unit_vector_2)
+        angle = math.degrees(np.arccos(dot_product))
+        return angle
+    else:
+        # print('pose_pred is None')
+        return None

microprograms/errors/splash_micro_program.py

@@ -0,0 +1,181 @@
+import sys, os, distutils.core
+
+# os.system('python -m pip install pyyaml==5.3.1')
+# dist = distutils.core.run_setup("./detectron2/setup.py")
+# temp = ' '.join([f"'{x}'" for x in dist.install_requires])
+# cmd = "python -m pip install {0}".format(temp)
+# os.system(cmd)
+sys.path.insert(0, os.path.abspath('./detectron2'))
+
+import detectron2
+import numpy as np
+import cv2
+
+# from detectron2.utils.logger import setup_logger
+# setup_logger()
+
+# # from detectron2.modeling import build_model
+# from detectron2 import model_zoo
+# from detectron2.engine import DefaultPredictor
+# from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+# from detectron2.data import MetadataCatalog, DatasetCatalog
+# from detectron2.checkpoint import DetectionCheckpointer
+# from detectron2.data.datasets import register_coco_instances
+
+# cfg = get_cfg()
+# cfg.OUTPUT_DIR = "./output/splash/"
+# # model = build_model(cfg)  # returns a torch.nn.Module
+# cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
+# cfg.DATASETS.TRAIN = ("splash_trains",)
+# cfg.DATASETS.TEST = ()
+# cfg.DATALOADER.NUM_WORKERS = 2
+# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
+# cfg.SOLVER.IMS_PER_BATCH = 2  # This is the real "batch size" commonly known to deep learning people
+# cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
+# cfg.SOLVER.MAX_ITER = 300    # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
+# cfg.SOLVER.STEPS = []        # do not decay learning rate
+# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # The "RoIHead batch size". 128 is faster, and good enough for this toy dataset (default: 512)
+# cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+# cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # path to the model we just trained
+# cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a custom testing threshold
+# predictor = DefaultPredictor(cfg)
+# register_coco_instances("splash_trains", {}, "./coco_annotations/splash/train.json", "./data/Splashes")
+# register_coco_instances("splash_vals", {}, "./coco_annotations/splash/val.json", "./data/Splashes")
+
+# from detectron2.utils.visualizer import ColorMode
+# splash_metadata = MetadataCatalog.get('splash_vals')
+# dataset_dicts = DatasetCatalog.get("splash_vals")
+
+# outputs_array = []
+# for d in dataset_dicts:
+#     im = cv2.imread(d["file_name"])
+#     outputs = predictor(im)
+#     outputs_array.append(outputs)  # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
+#     v = Visualizer(im[:, :, ::-1],
+#                    metadata=splash_metadata,
+#                    scale=0.5,
+#                    instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models
+#     )
+#     out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+#     img = out.get_image()[:, :, ::-1]
+#     filename = os.path.join("./output", d["file_name"][2:])
+#     if not cv2.imwrite(filename, img):
+#       print('no image written')
+
+import torch
+# SPLASH MICRO PROGRAM
+# find percentage of "True" in pred_masks, lower the percentage the better. get from outputs
+# may need to calibrate to points
+
+def get_splash_pred_mask(output):
+  pred_classes = output['instances'].pred_classes.cpu().numpy()
+  splashes = np.where(pred_classes == 0)[0]
+  scores = output['instances'].scores[splashes]
+  if len(scores) == 0:
+    return None
+  pred_masks = output['instances'].pred_masks[splashes]
+  max_instance = torch.argmax(scores)
+  # pred_mask = pred_masks[max_instance] # splash instance with highest confidence
+  pred_mask = np.array(pred_masks[max_instance].cpu()) 
+  return pred_mask
+
+# function that finds the splash instance with the highest percent confidence
+# and returns the
+def splash_area_percentage(output, pred_mask=None):
+  if pred_mask is None:
+    return
+  # loops over pixels to get sum of splash pixels
+  totalSum = 0
+  for j in range(len(pred_mask)):
+    totalSum += pred_mask[j].sum()
+  # return percentage of image that is splash
+  return totalSum/(len(pred_mask) * len(pred_mask[0]))
+
+
+# loops over each image
+# i = 0
+# for output in outputs_array:
+#   print(dataset_dicts[i]["file_name"])
+#   # print(output)
+#   print(splash_area_percentage(output))
+#   i+=1
+
+
+# TODO: run splash micro program on one diving clip
+# plot splash area percentage and save image
+
+import matplotlib.pyplot as plt
+from models.detectron2.splash_detector_setup import get_splash_detector
+
+def get_splash_from_one_frame(filepath, im=None, predictor=None, visualize=False, dive_folder_num=""):
+    if predictor is None:
+      predictor=get_splash_detector()
+    if im is None:
+      im = cv2.imread(filepath)
+    outputs = predictor(im)
+    pred_mask = get_splash_pred_mask(outputs)
+    area = splash_area_percentage(outputs, pred_mask=pred_mask)
+    if area is None:
+        # print("no splash detected in", filepath)
+        return None, None
+    if visualize:
+        pred_boxes = outputs['instances'].pred_boxes
+        print("pred_boxes", pred_boxes)
+        for box in pred_boxes:
+          image = cv2.rectangle(im, (int(box[0]),int(box[1])), (int(box[2]),int(box[3])), color=(0, 0, 255), thickness=2)
+          out_folder= "./output/data/splash/{}".format(dive_folder_num)
+          if not os.path.exists(out_folder):
+              os.makedirs(out_folder)
+          filename = os.path.join(out_folder, filepath.split('/')[-1])
+          if not cv2.imwrite(filename, image):
+              print('no image written to', filename)
+          break
+        
+    return area.tolist(), pred_mask
+
+# outputs_array2 = []
+# directory = "./MTL-AQA/"
+# file_names = os.listdir(directory)
+# for file_name in file_names:
+#     if file_name[:10] == "img_01_06_":
+#         path = os.path.join(directory, file_name)
+#         # im = caffe.io.load_image(path)
+#         print("PATH IM_01_06:", path)
+#         im = cv2.imread(path)
+#         outputs = predictor(im)
+#         outputs_array2.append(outputs)  # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
+#         v = Visualizer(im[:, :, ::-1],
+#                       metadata=splash_metadata,
+#                       scale=0.5,
+#                       instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models
+#         )
+#         out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+#         img = out.get_image()[:, :, ::-1]
+#         filename = os.path.join("./output/data/img_01_06/", file_name)
+#         if not cv2.imwrite(filename, img):
+#           print('no image written')
+
+# i = 0
+# splash_area = []
+# for output in outputs_array2:
+#   # print(output)
+#   area = splash_area_percentage(output)
+#   if area is None:
+#     splash_area.append(0)
+#   else:
+#     print(area.cpu().data.item())
+#     splash_area.append(area.cpu().data.item())
+#   # print(splash_area_percentage(output))
+#   i+=1
+
+# print(range(i))
+# print(splash_area)
+# plt.plot(range(i), splash_area)
+# plt.savefig('./output/data/img_01_06/img_01_06_splash_graph.png')
+
+# from detectron2.evaluation import COCOEvaluator, inference_on_dataset
+# from detectron2.data import build_detection_test_loader
+# evaluator = COCOEvaluator("splash_vals", output_dir="./output")
+# val_loader = build_detection_test_loader(cfg, "splash_vals")
+# print(inference_on_dataset(predictor.model, val_loader, evaluator))

microprograms/temporal_segmentation/entry.py

@@ -0,0 +1,30 @@
+from microprograms.errors.splash_micro_program import get_splash_from_one_frame
+from microprograms.errors.angles_micro_programs import applyPositionTightnessError
+import numpy as np
+
+def entry_microprogram_one_frame(filepath, above_board, on_board, pose_pred, expected_twists, petal_count, expected_som, half_som_count, frame=None, splash_detector=None,  visualize=False, dive_folder_num=None):
+    if above_board:
+        return 0
+    if on_board:
+        return 0
+    splash = get_splash_from_one_frame(filepath, im=frame, predictor=splash_detector, visualize=visualize, dive_folder_num=dive_folder_num)
+    if splash:
+        return 1
+    # if completed with somersaults, we know we're in entry phase
+    if not expected_som > half_som_count:
+        return 1
+    if expected_twists > petal_count or expected_som > half_som_count:
+        return 0
+    return 1
+    
+    # if pose_pred is None:
+    #     return 1
+    # else:
+    #     angle = applyPositionTightnessError(filepath, pose_pred)
+    #     if angle > 90:
+    #         return 1
+    # if splash is not None:
+    #     return 1
+    # else:
+    #     return 0
+    

microprograms/temporal_segmentation/somersault.py

@@ -0,0 +1,20 @@
+from microprograms.errors.angles_micro_programs import applyPositionTightnessError
+# from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+
+def somersault_microprogram_one_frame(filepath, on_board, expected_som, half_som_count, expected_twists, petal_count, pose_pred=None, diver_detector=None, pose_model=None):
+    if on_board:
+        return 0
+    if expected_som <= half_som_count:
+        return 0
+    # if not done with som or twists, need to determine if som or twist
+    angle = applyPositionTightnessError(filepath, pose_pred, diver_detector=diver_detector, pose_model=pose_model)
+    if angle is None:
+        return 0
+    # if not done with som but done with twists
+    if expected_som > half_som_count and expected_twists <= petal_count:
+        return 1
+    # print("angle:", angle)
+    if angle <= 80:
+        return 1
+    else:
+        return 0

microprograms/temporal_segmentation/start_takeoff.py

@@ -0,0 +1,22 @@
+from microprograms.errors.angles_micro_programs import applyPositionTightnessError
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+from microprograms.errors.distance_from_springboard_micro_program import board_end
+
+from microprograms.errors.distance_from_springboard_micro_program import find_which_side_board_on
+from models.detectron2.platform_detector_setup import get_platform_detector
+import numpy as np
+import cv2
+
+def takeoff_microprogram_one_frame(filepath, above_board, on_board, pose_pred=None):
+    if not above_board:
+        return 0
+    if on_board:
+        return 1
+    return 0
+    # angle = applyPositionTightnessError(filepath, pose_pred)
+    # if angle is None:
+    #     return 0
+    # if angle > 90:
+    #     return 1
+    # else:
+    #     return 0

microprograms/temporal_segmentation/twist.py

@@ -0,0 +1,21 @@
+from microprograms.errors.angles_micro_programs import applyPositionTightnessError
+from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
+from microprograms.errors.distance_from_springboard_micro_program import board_end
+
+from microprograms.errors.distance_from_springboard_micro_program import find_which_side_board_on
+from models.detectron2.platform_detector_setup import get_platform_detector
+import cv2
+import numpy as np
+
+def twist_microprogram_one_frame(filepath, on_board, expected_twists, petal_count, expected_som, half_som_count, pose_pred=None, diver_detector=None, pose_model=None):
+    if on_board:
+        return 0
+    if expected_twists <= petal_count or expected_som <= half_som_count:
+        return 0
+    angle = applyPositionTightnessError(filepath, pose_pred=pose_pred, diver_detector=diver_detector, pose_model=pose_model)
+    if angle is None:
+        return 0
+    if angle > 80:
+        return 1
+    else:
+        return 0