def get_default_instance_segmentation_config(use_direction_classes,
training_set_name="train",
test_set_name="val",
max_iter=1000,
threshold=0.7):
cfg = get_cfg()
cfg.DATA_TRANSFORMATIONS = CN()
cfg.DATA_TRANSFORMATIONS.ROTATION = True
cfg.DATA_TRANSFORMATIONS.FLIP = True
cfg.DATA_TRANSFORMATIONS.RESIZE_FACTOR = 1
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
) # only segmentation and bounding boxes
cfg.TYPE = "mask_rcnn_R_50_FPN_3x"
assert cfg.TYPE in cfg.MODEL.WEIGHTS, "miss match between type and weights"
# General
cfg.SEED = 42
# Data set
cfg.DATASETS.TRAIN = (get_name_with_prefix(training_set_name,
use_direction_classes), )
cfg.DATASETS.TEST = (get_name_with_prefix(test_set_name,
use_direction_classes),
) # TODO can I add test set also here?
cfg.DATASETS.USE_DIRECTION_CLASSES = use_direction_classes
cfg.DATALOADER.NUM_WORKERS = 0
# validation
cfg.TEST.EVAL_PERIOD = 60
# Solver
cfg.SOLVER.MAX_ITER = max_iter
cfg.SOLVER.IMS_PER_BATCH = 5
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.WARMUP_ITERS = int(0.5 * cfg.SOLVER.MAX_ITER)
cfg.SOLVER.WARMUP_FACTOR = 1.0 / (cfg.SOLVER.WARMUP_ITERS + 1)
cfg.SOLVER.WEIGHT_DECAY_NORM = 0.0
# Model
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 if not use_direction_classes else len(
DIRECTION_CLASSES)
cfg.MODEL.BACKBONE.FREEZE_AT = 2
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
update_output_dir_path(cfg)
return cfg
def main(image_dir,
project,
crop_size,
batch_size,
iterations,
validation_split,
backbone,
output_dir,
learning_rate,
device,
tb,
tb_port,
resume=None):
if tb:
print(f'Tensorboard URL: {launch_tb(output_dir, tb_port)}')
DatasetCatalog.register("train",
lambda: get_balloon_dicts(image_dir, project))
MetadataCatalog.get("train").set(thing_classes=["balloon"])
DatasetCatalog.register("val",
lambda: get_balloon_dicts(image_dir, project))
MetadataCatalog.get("val").set(thing_classes=["balloon"])
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("train", )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
if resume:
cfg.MODEL.WEIGHTS = resume
else:
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
f"COCO-InstanceSegmentation/mask_rcnn_{backbone}.yaml"
) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = batch_size
cfg.SOLVER.BASE_LR = learning_rate # pick a good LR
cfg.SOLVER.MAX_ITER = iterations # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.MODEL.DEVICE = device
cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS = False
cfg.OUTPUT_DIR = output_dir
cfg.INPUT.CROP.SIZE = crop_size
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = Trainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def get_cfg(model_weights_path: Path = None,
output_path: Path = None,
debug: bool = True,
num_input_channels: int = 1):
cfg = detectron_get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
# cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml"))
if model_weights_path is None:
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
else:
cfg.MODEL.WEIGHTS = str(model_weights_path)
cfg.MODEL.RESNETS.NORM = "BN"
cfg.MODEL.BACKBONE.FREEZE_AT = 0
cfg.OUTPUT_DIR = str(
output_path) if output_path is not None else './output'
Path(cfg.OUTPUT_DIR).mkdir(exist_ok=True)
cfg.DATALOADER.NUM_WORKERS = 0 if debug else 6
cfg.SOLVER.WARMUP_FACTOR = 1.0 / 500
cfg.SOLVER.WARMUP_ITERS = 500 # a warm up is necessary to avoid diverging training while keeping the goal learning rate as high as possible
cfg.SOLVER.IMS_PER_BATCH = 16 if not debug else 8
cfg.SOLVER.BASE_LR = 0.0005 # pick a good LR
cfg.SOLVER.MAX_ITER = 80000
cfg.SOLVER.STEPS = (40000, 60000, 70000)
cfg.SOLVER.GAMMA = 0.5 # after each milestone in SOLVER.STEPS gets reached, the learning rate gets scaled by Gamma.
cfg.SOLVER.CHECKPOINT_PERIOD = 50 if debug else 3000 #5000
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.TEST.EVAL_PERIOD = 30 if debug else 3000
cfg.INPUT.MASK_FORMAT = "bitmask"
cfg.INPUT.FORMAT = "D" * num_input_channels
cfg.MODEL.PIXEL_MEAN = [0.5] * num_input_channels
cfg.MODEL.PIXEL_STD = [1.0] * num_input_channels
cfg.MIN_AREA = 100
cfg.DATASETS.TRAIN = ("val", ) if debug else ("train", )
cfg.DATASETS.TEST = ("val", )
cfg.DEBUG = debug
return cfg
def main():
args = parser.parse_args()
### regist datasethasattr
register_coco_instances(args.dataset, {}, args.label, args.file)
# detectron2.data.datasets.load_coco_json(args.label, args.file, "ACID_dataset") # this will set thing_classes = ["excavator", "dump_truck", "cement_truck"]
## for keypoint training
# MetadataCatalog.get("ACID_dataset").keypoint_names = ['body_end_x', 'body_end_y', 'body_end_v', 'cab_boom_x', 'cab_boom_y', 'cab_boom_v', 'boom_arm_x', 'boom_arm_y', 'boom_arm_v',
# 'arm_bucket_x', 'arm_bucket_y', 'arm_bucket_v', 'bucket_end_left_x', 'bucket_end_left_y', 'bucket_end_left_v', 'bucket_end_right_x', 'bucket_end_right_y', 'bucket_end_right_v']
# MetadataCatalog.get("ACID_dataset").keypoint_flip_map = []
# MetadataCatalog.get("ACID_dataset").keypoint_connection_rules = []
## end for keypoint training
### set metadata
ACID_meta = MetadataCatalog.get(args.dataset)
dataset_dicts = DatasetCatalog.get(args.dataset)
### verify the data loading is correct
# visualizeDataset(dataset_dicts, ACID_meta)
### train model
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(args.model))
cfg.DATASETS.TRAIN = (args.dataset, )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
args.model) # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.02 #0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = args.iter # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 # 128 # default: 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = args.num_class # excavator, dump_truck, cement_truck
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
# torch.cuda.empty_cache()
trainer.train()
def new_model_cfg():
cfg = get_cfg()
model = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
cfg.merge_from_file(model_zoo.get_config_file(model))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model)
cfg.MODEL.MASK_ON = False
cfg.INPUT.RANDOM_FLIP = "none"
cfg.OUTPUT_DIR = "output"
cfg.DATALOADER.NUM_WORKERS = 4
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.CHECKPOINT_PERIOD = 2000
cfg.SOLVER.STEPS = (10000, 21000, 35000, 55000, 75000)
cfg.SOLVER.MAX_ITER = 150000
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.WEIGHT_DECAY = 0.001
cfg.TEST.EVAL_PERIOD = 2000
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS = True
return cfg
def create_cfg(self):
"""
Creates configuration file for the model.
:return:
"""
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
# Passing the Train and Validation sets
cfg.DATASETS.TRAIN = ("smilesdetect_train", )
cfg.DATASETS.TEST = ("smilesdetect_val", )
cfg.OUTPUT_DIR = self.base_path + '/trained_models'
cfg.INPUT.FORMAT = self.input_format
# Number of data loading threads
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml")
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(self.unique_labels)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
return cfg
def define_model(train_dts, val_dts, device, model, lr, iterations,
batch_size):
cfg = get_cfg()
cfg.OUTPUT_DIR = "output-training-" + datetime.datetime.now().strftime(
"%d-%m-%Y-(%H:%M:%S)")
cfg.merge_from_file(model_zoo.get_config_file(model))
cfg.MODEL.DEVICE = device
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model)
cfg.DATASETS.TRAIN = (train_dts, )
cfg.DATASETS.TEST = (val_dts, )
cfg.DATALOADER.NUM_WORKERS = 4
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = lr
cfg.SOLVER.MAX_ITER = iterations
cfg.SOLVER.STEPS = [] # do not decay learning rate
cfg.TEST.EVAL_PERIOD = 50
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = batch_size
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(grouped_classes_dict.keys())
return cfg
def predict_surface(img):
for d in ["train", "val"]:
DatasetCatalog.clear()
DatasetCatalog.register("surface_" + d,
lambda d=d: get_surface_dicts("surface/" + d))
MetadataCatalog.get("surface_" + d).set(thing_classes=["surface"])
surface_metadata = MetadataCatalog.get("surface_train")
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("surface_train", )
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")
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.SOLVER.BASE_LR = 0.0025
cfg.SOLVER.MAX_ITER = 1000
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.OUTPUT_DIR = "./drive/MyDrive/surface"
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3 # set a custom testing threshold
predictor = DefaultPredictor(cfg)
dataset_dicts = "surface_img/val"
#files=sorted(glob.glob("balloon/val/*.jpg"))
teller = 1
data = []
middle = [img[0].shape[0] / 2, img[0].shape[1] / 2]
background = backgroundsubtraction(img)
cv2.imwrite("background.jpg", background)
for teller, k in enumerate(img): #go through all images
minimum = None
predictor = DefaultPredictor(cfg)
outputs = predictor(k) #predict surface of image k
"""
v = Visualizer(k[:, :, ::-1], metadata=surface_metadata, scale=0.5) #visualization of predicted surface
out=v.draw_instance_predictions(outputs["instances"].to("cpu")) #...
v=out.get_image()[:, :, ::-1] #...
cv2.imwrite("./drive/MyDrive/wkvideo/surface/"+str(teller)+".jpg",v) #...
"""
#maskoutput=outputs['instances'].pred_masks.to("cpu")[0][:2] #next paragraph: check for multiple detections and make a
#decision if the predicted surface is part of road where the cyclist ride on
maskoutput = 0
indexen = []
y = []
x = []
coordinaten = []
prev_x_min = 0
prev_x_max = k.shape[1]
if len(
outputs['instances'].pred_boxes
) == 0: #if only one road detection -> take detection as "finish road"
maskoutput = np.zeros((k.shape[0], k.shape[1]),
np.uint8) #make new mask with certain shape
else:
for index, k in enumerate(
outputs['instances'].pred_boxes.to("cpu")
): #if multiple road detections-> do some checks about position of detection(knowledge that "finish road" is mostly in middle of frame)
#go through all predictions and take boundingboxescoordinates as variable k
coordinates = k.numpy(
) #transform k from tensor object to numpy
middle = coordinates[2] - coordinates[
0] #calculate the center position of the bounding box in x direction
if middle >= prev_x_min and middle <= prev_x_max: #check if center position is between two boundaries( prev frame boundaries of bounding box position)
y.append(
coordinates[3] -
coordinates[1]) #add center of boundingbox to list
x.append(coordinates[2] - coordinates[0])
indexen.append(
index) #add index of detected boundingbox to list
coordinaten.append(
coordinates) #add coordinates of boundingbox to list
best_ind = 0 #in next paragraph: if multiple surfaces where detected: do some logic,
#to find the best matching detected surface, using the position of the last detected surface of previous image
if len(indexen) > 1:
best = None
lastone = False
for d, k in enumerate(indexen[:len(indexen) - 1]):
if abs(x[d] - x[d + 1]) > ((prev_x_max - prev_x_min) / 2):
if d == len(indexen) - 1:
lastone = True
dist = abs(x[d] - (prev_x_max - prev_x_min) / 2)
if best == None or dist < best:
best = dist
best_ind = d
if lastone:
d = len(indexen) - 1
dist = abs(x[d] - (prev_x_max - prev_x_min) / 2)
if best == None or dist < best:
best = dist
best_ind = d
indexen = [best_ind]
prev_x_min = coordinaten[best_ind][0]
prev_x_max = coordinaten[best_ind][2]
for index, k in enumerate(
outputs['instances'].pred_masks.to("cpu").numpy()):
if indexen.count(index) == 1:
maskoutput += k
maskoutput = maskoutput * 255
kernel = np.ones((9, 1), np.uint8)
maskoutput = maskoutput.astype(np.uint8)
maskoutput = cv2.dilate(maskoutput, kernel, iterations=4)
maskoutput += background
maskoutput = maskoutput.astype(np.uint8)
mask = np.ones((k.shape[0], k.shape[1]), dtype=np.uint8)
img_res = cv2.bitwise_and(mask, mask, mask=maskoutput)
data.append(img_res)
del (indexen)
del (y)
del (x)
DatasetCatalog.clear()
return data
DatasetCatalog.register("renner_" + d,
lambda d=d: get_balloon_dicts("surface_img/" + d))
MetadataCatalog.get("renner_" + d).set(thing_classes=["renner"])
balloon_metadata = MetadataCatalog.get("renner_train")
#dataset_dicts = get_balloon_dicts("surface_img/train")
cfg = get_cfg()
cfg.MODEL.DEVICE = 'cpu'
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("renner_train", )
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 = 1
cfg.SOLVER.BASE_LR = 0.0025 #0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 1000 # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon)
cfg.OUTPUT_DIR = "./drive/MyDrive/renner"
# cfg already contains everything we've set previously. Now we changed it a little bit for inference:
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.5 #0.7 # set a custom testing threshold
predictor = DefaultPredictor(cfg)
#dataset_dicts = get_balloon_dicts("balloon/train")
dataset_dicts = "surface_img/val"
from detectron2.modeling import build_model
register_coco_instances('my_dataset_train', {}, './data/train.json',
'./data/train_images')
register_coco_instances('my_dataset_val', {}, './data/val.json',
'./data/train_images')
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
'COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml'))
cfg.DATASETS.TRAIN = ('my_dataset_train', )
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 4
# Let training initialize from model zoo
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
'COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 1000
# 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 20
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
# Inference should use the config with parameters that are used in training
# cfg now already contains everything we've set previously. We changed it a little bit for inference:
cfg.MODEL.WEIGHTS = os.path.join(
import glob
import multiprocessing as mp
import os
import time
import cv2
from detectron2.config import get_cfg # returns a copy of the default config
from detectron2.engine import DefaultTrainer
from detectron2.model_zoo import model_zoo
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"
) # Let training initialize from model zoo
# resnext ref: https://medium.com/@14prakash/understanding-and-implementing-architectures-of-resnet-and-resnext-for-state-of-the-art-image-cc5d0adf648e
cfg.MODEL.RESNETS.NUM_GROUPS = 32 # 1 ==> ResNet; > 1 ==> ResNeXt
cfg.MODEL.RESNETS.WIDTH_PER_GROUP = 8
cfg.MODEL.RESNETS.DEPTH = 101
cfg.MODEL.DEVICE = 'cpu'
#dataset default is coco
#cfg.DATASETS.TRAIN = ("",)
#cfg.DATASETS.TEST = ("",)
cfg.INPUT.MIN_SIZE_TRAIN = (100, )
cfg.INPUT.MAX_SIZE_TRAIN = 166
cfg.INPUT.MIN_SIZE_TEST = (100, )
cfg.INPUT.MAX_SIZE_TEST = 166
def run_colmap(video_data, frame_data_path, mask_data_path,
colmap_workspace_path, logger):
cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
def get_mask(frame):
segmentation_outputs = predictor(frame)
masks = segmentation_outputs["instances"].pred_masks.cpu().numpy()
pred_classes = segmentation_outputs["instances"].pred_classes.cpu(
).numpy()
filtered_masks = []
for mask, pred_class in zip(masks, pred_classes):
# The class zero is "people".
if pred_class == 0:
filtered_masks.append(mask)
filtered_masks = np.array(filtered_masks)
flattened_mask = np.zeros(shape=(masks.shape[-2:]), dtype=bool)
for mask in filtered_masks:
flattened_mask |= mask
flattened_mask = np.stack((flattened_mask, ) * 3, axis=-1)
return flattened_mask
if not os.path.exists(frame_data_path):
os.makedirs(frame_data_path)
logger.log(
"Created folder for frame data at {}.".format(frame_data_path))
os.makedirs(mask_data_path)
logger.log(
"Created folder for mask data at {}.".format(mask_data_path))
for frame_i, frame in enumerate(video_data):
frame_filename = "{:04d}.png".format(frame_i + 1)
# If the video data is in the RGB format (as opposed to the BGR format), then it must be converted to BGR
# before being written to disk since that is the format OpenCV uses.
if video_data.is_rgb:
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
frame_path = os.path.join(frame_data_path, frame_filename)
cv2.imwrite(frame_path, frame)
mask_file_name = f"{frame_filename}.png"
mask_path = os.path.join(mask_data_path, mask_file_name)
instance_mask = get_mask(frame)
# Need to invert mask so that the masked areas are 'False' (zero) instead of True (one) and so COLMAP only
# masks out people.
instance_mask = ~instance_mask
cv2.imwrite(mask_path, instance_mask.astype(np.uint8))
logger.log(
f"Wrote frame and mask {frame_i + 1:,d}/{len(video_data):,d} to {frame_path} and {mask_path}.\r",
end="")
print()
logger.log(
f"Wrote frame data for COLMAP to {frame_data_path} and dynamic object masks to {mask_data_path}."
)
else:
logger.log(f"Found frame data at {frame_data_path}.")
if not os.path.exists(colmap_workspace_path):
os.makedirs(colmap_workspace_path)
logger.log(
"Created workspace folder at {}.".format(colmap_workspace_path))
logger.log("Running COLMAP reconstruction. This may take a while...")
colmap_process = subprocess.run([
'colmap', 'automatic_reconstructor', '--image_path',
frame_data_path, '--mask_path', mask_data_path, '--workspace_path',
colmap_workspace_path, '--single_camera', '1', '--quality', 'low',
'--data_type', 'video', '--camera_model', 'SIMPLE_PINHOLE',
'--sparse', '1', '--dense', '0'
])
if colmap_process.returncode != 0:
raise RuntimeError(
"COLMAP exited with the non-zero return code {}.".format(
colmap_process.returncode))
else:
logger.log("COLMAP finished processing the video.")
else:
logger.log(
"Found COLMAP reconstruction workspace folder at {}.".format(
colmap_workspace_path))