Spaces:

X-NS
/

NSAQA

Running

File size: 12,497 Bytes

97a245c

# 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')