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