def crop_image_mask_and_save(image, outputs, result_images_dir, vis):
masks = np.asarray(outputs["instances"].pred_masks)
masks = [
GenericMask(x, vis.output.height, vis.output.width) for x in masks
]
areas = np.asarray([x.area() for x in masks])
sorted_idxs = np.argsort(-areas).tolist()
masks = [masks[idx] for idx in sorted_idxs]
boxes = [outputs["instances"].pred_boxes[idx] for idx in sorted_idxs]
for i, box in zip(range(len(masks)), boxes):
background = np.ones(image.shape, dtype=np.uint8)
background.fill(255)
for j, segment in enumerate(masks[i].polygons):
item_mask = np.array([segment.reshape(-1, 2)], dtype=np.int32)
cv2.fillPoly(background, item_mask, 0)
masked_image = cv2.bitwise_or(image, background)
x1, y1, x2, y2 = [round(cord) for cord in box.tensor.tolist()[0]]
crop_img = masked_image[y1:y2, x1:x2]
cropped_image_path = Path(result_images_dir,
f"{uuid.uuid4()}.jpg").as_posix()
cv2.imwrite(cropped_image_path, crop_img)
def format_results(self, class_names):
"""
Format results so they can be used by overlay_instances function
"""
predictions = self.outputs['instances']
boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has("pred_classes") else None
labels = None
if classes is not None and class_names is not None and len(class_names) > 1:
labels = [class_names[i] for i in classes]
if scores is not None:
if labels is None:
labels = ["{:.0f}%".format(s * 100) for s in scores]
else:
labels = ["{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores)]
masks = predictions.pred_masks.cpu().numpy()
masks = [GenericMask(x, v.output.height, v.output.width) for x in masks]
boxes_list = boxes.tensor.tolist()
scores_list = scores.tolist()
class_list = classes.tolist()
for i in range(len(scores_list)):
boxes_list[i].append(scores_list[i])
boxes_list[i].append(class_list[i])
boxes_list = np.array(boxes_list)
return (masks, boxes, boxes_list, labels, scores_list, class_list, classes)
def return_attributes(self, path):
im = cv2.imread(path)
output = self.predictor(im)
print('STATUS: Detectron2 job done.')
self.SCORE_THRESHOLD = 0.7
AREA_FRACTION_THRESHOLD = 0.1
# storing attributes
predictions = output["instances"].to("cpu")
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has("pred_classes") else None
if predictions.has("pred_masks"):
masks = [GenericMask(x, im.shape[0], im.shape[1]) for x in np.asarray(predictions.pred_masks)]
else:
masks = None
humans = (classes==0).nonzero().reshape(-1).numpy()
if(humans.shape[0]==0):
return None
temp = []
for i in range(len(humans)):
if(scores[humans[i]]>self.SCORE_THRESHOLD and masks[i].area()/(im.shape[0] * im.shape[1])>self.AREA_FRACTION_THRESHOLD):
# if(True):
temp.append(humans[i])
humans = temp
human_masks = []
for i in range(len(masks)):
if(i in humans):
human_masks.append(masks[i])
return im, human_masks
def draw_instance_predictions(self, predictions):
"""
Draw instance-level prediction results on an image.
Args:
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
Returns:
output (VisImage): image object with visualizations.
"""
boxes = predictions.pred_boxes if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
contacts = predictions.pred_cats if predictions.has(
"pred_cats") else None
scores = contacts * scores.reshape(-1, 1)
labels, dis_colors = _create_custom_text_labels(
classes, scores, "hand", self.scores_thresh)
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
if predictions.has("pred_masks"):
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, self.output.height, self.output.width)
for x in masks
]
else:
masks = None
if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get(
"thing_colors"):
colors = [
self._jitter([x / 255 for x in self.metadata.thing_colors[c]])
for c in classes
]
alpha = 0.8
else:
colors = None
alpha = 0.5
if self._instance_mode == ColorMode.IMAGE_BW:
self.output.img = self._create_grayscale_image((
predictions.pred_masks.any(dim=0) > 0
).numpy() if predictions.has("pred_masks") else None)
alpha = 0.3
self.overlay_instances(
masks=masks,
boxes=None,
labels=labels,
keypoints=keypoints,
assigned_colors=dis_colors,
alpha=alpha,
)
return self.output
def draw_instance_predictions(self, predictions):
"""
Draw instance-level prediction results on an image.
Args:
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
Returns:
output (VisImage): image object with visualizations.
"""
#boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None
boxes = None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
#labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None))
labels = None
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
if predictions.has("pred_masks"):
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, self.output.height, self.output.width)
for x in masks
]
else:
masks = None
colors = []
alpha = 0.8
for c in classes:
color = (0, 1, 1, 1) #Green
if c == 0:
color = (1, 1, 0, 1) #Yellow Line
if c == 1:
color = (0.8, 0.8, 1, 1) #White Line
if c == 4:
color = (1, 0, 0, 1) #Red line
if c == 2:
color = (0, 1, 0, 1) #Obstacle
colors.append(color)
if self._instance_mode == ColorMode.IMAGE_BW:
self.output.img = self._create_grayscale_image((
predictions.pred_masks.any(dim=0) > 0
).numpy() if predictions.has("pred_masks") else None)
alpha = 0.3
self.overlay_instances(
masks=masks,
boxes=boxes,
labels=labels,
keypoints=keypoints,
assigned_colors=colors,
alpha=alpha,
)
return self.output
def predict_image(im):
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)
outputs = predictor(im)
'''
look at the outputs. See
https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
for specification
'''
print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)
# We can use `Visualizer` to draw the predictions on the image.
'''
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.2)
'''
v = Visualizer(np.zeros_like(im),
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.0)
im_mask = np.zeros_like(im, np.uint8)
if outputs["instances"].has("pred_masks"):
masks = np.asarray(outputs["instances"].to("cpu").pred_masks)
value_statis(masks, 'massk1')
print('masks1: {}'.format(type(masks)))
for idx_i in range(masks.shape[0]):
im_mask[masks[idx_i, :, :]] = idx_i * 2 + 10
masks = [
GenericMask(x, v.output.height, v.output.width) for x in masks
]
print('masks2: {}'.format(type(masks)))
num_instance = len(v._convert_masks(masks))
print(' -> num_instances: {}'.format(num_instance))
out = v.overlay_instances(masks=masks)
else:
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
value_statis(out.get_image(), 'out.get_image()')
value_statis(im_mask, 'im_mask')
# return out.get_image()
return im_mask
def draw_instance_predictions(self, predictions, track_ids):
"""
Draw instance-level prediction results on an image.
Args:
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
Returns:
output (VisImage): image object with visualizations.
"""
boxes = predictions.pred_boxes if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
labels = _create_text_labels(classes, scores,
self.metadata.get("thing_classes", None))
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
if predictions.has("pred_masks"):
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, self.output.height, self.output.width)
for x in masks
]
else:
masks = None
# set the color according to the track ids
colors = [cm.tab20(id_) for id_ in track_ids]
alpha = 0.6
labels = [
f'Track {id_} {label}' for label, id_ in zip(labels, track_ids)
]
# increase font size
if self._default_font_size < 20: self._default_font_size *= 1.3
if self._instance_mode == ColorMode.IMAGE_BW:
assert predictions.has(
"pred_masks"), "ColorMode.IMAGE_BW requires segmentations"
self.output.img = self._create_grayscale_image(
(predictions.pred_masks.any(dim=0) > 0).numpy())
self.overlay_instances(
masks=masks,
boxes=boxes,
labels=labels,
keypoints=keypoints,
assigned_colors=colors,
alpha=alpha,
)
return self.output
def create_annotated_scan(original, mask, output_dir, suffix):
original = original.copy()
mask = GenericMask(mask, *mask.shape)
polygons = [poly.reshape(-1, 2) for poly in mask.polygons]
cv2.polylines(original,
polygons,
isClosed=True,
color=(0, 255, 0),
thickness=3)
cv2.imwrite(os.path.join(output_dir, f'{suffix}.jpg'), original)
return cv2.imread(os.path.join(output_dir, f'{suffix}.jpg'))
def format_predictions(predictions, height, width):
masks = np.asarray(predictions.pred_masks)
masks = [GenericMask(mask, height, width) for mask in masks]
# boxes = predictions.pred_boxes if predictions.has("pred_boxes") else [mask.bbox() for mask in masks]
classes = predictions.pred_classes if predictions.has(
"pred_classes") else [None for _ in masks]
areas = [mask.area() for mask in masks]
return [(reshape_and_close_poly(mask.polygons[0]), area, cls)
for mask, area, cls in zip(masks, classes, areas)
if len(mask.polygons) > 0]
def predOneImg(predictor, imgNameFull, annoResName, visRoot, thing_classes, divNum=2, isVis = 0):
img = cv2.imread(imgNameFull)
hImg, wImg = img.shape[0], img.shape[1]
stepH, stepW = int(hImg/divNum), int(wImg/divNum)
imgName = str.split(imgNameFull, '/')[-1]
resDict = {'filename':imgName, 'labels':[]}
for ii in range(len(thing_classes)):
resDict['labels'].append({'name':thing_classes[ii], 'annotations':[]})
if isVis == 1:
vPoly = Visualizer(img[:, :, ::-1],
scale=1.0,
instance_mode=ColorMode.IMAGE_BW
)
segmentCount = 0
for ii in range(divNum):
startH = ii * stepH
endH = hImg if ii == divNum - 1 else stepH*(ii + 1) -1
for jj in range(divNum):
print(ii, jj)
startW = jj * stepW
endW = wImg if jj == divNum - 1 else stepW * (jj + 1) -1
tmpImg = img[startH:endH+1, startW:endW+1, :]
outputs = predictor(tmpImg)
res = outputs["instances"].to("cpu")
height = res.image_size[0]
width = res.image_size[1]
masks = [GenericMask(x.numpy(), height, width).mask_to_polygons(x.numpy()) for x in res.pred_masks]
# masks1 = [x[0][0].reshape(-1, 2)+[startW, startH] for x in masks]
labels = [int(x.numpy()) for x in res.pred_classes]
colorMaps = [[0., 0., 1.], [1., 0., 0.], [1.0, 0.3, .5]]
for ll in range(len(labels)):
tmpLabel = thing_classes[labels[ll]]
if len(masks[ll][0]) == 0:
continue
tmpsegmentNP = masks[ll][0][0].reshape(-1, 2)+[startW, startH]
tmpsegmentRDP = rdp(tmpsegmentNP, epsilon=6)
tmpsegment = tmpsegmentRDP.reshape(-1, 1).squeeze().tolist()
tmpInst = {'id':segmentCount, 'segmentation':tmpsegment}
for cc in range(len(resDict['labels'])):
if resDict['labels'][cc]['name'] == tmpLabel:
resDict['labels'][cc]['annotations'].append(tmpInst)
if isVis == 1:
vPoly.output = vPoly.draw_polygon(tmpsegmentRDP, colorMaps[cc], alpha=0.3)
# plt.imshow(vPoly.output.get_image()[:, :, ::-1])
# plt.show()
segmentCount += 1
with open(annoResName, 'w') as fp:
json.dump(resDict, fp)
if isVis == 1:
picName = os.path.join(visRoot, imgName)
plt.imsave(picName, vPoly.output.get_image()[:, :, ::-1])
def unwrap(predictions, imageShape):
boxes = predictions.pred_boxes.tensor.numpy() if predictions.has(
"pred_boxes") else None
scores = predictions.scores.numpy() if predictions.has("scores") else None
classes = predictions.pred_classes.numpy() if predictions.has(
"pred_classes") else None
if predictions.has("pred_masks"):
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, imageShape[0], imageShape[1]).mask for x in masks
]
else:
masks = None
return boxes, masks, scores, classes
def to_polygons(predictions):
if type(predictions) == np.ndarray:
masks = [
GenericMask(predictions, predictions.shape[0],
predictions.shape[1])
]
else:
boxes = predictions.pred_boxes if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
if not predictions.has('pred_masks'):
return None
masks = np.asarray(predictions.pred_masks)
masks = [GenericMask(x, x.shape[0], x.shape[1]) for x in masks]
polygons = [
poly.reshape(-1, 2) for mask in masks for poly in mask.polygons
]
return polygons
def extract_predictions(predictions):
if not predictions.has('pred_masks'):
return None, None
masks = np.asarray(predictions.pred_masks)
if masks.shape[0] == 0:
return None, None
mask = np.zeros_like(masks[0, :, :])
for m in masks:
mask |= m
masks = [GenericMask(m, m.shape[0], m.shape[1]) for m in masks]
polygons = [
poly.reshape(-1, 2) for mask in masks for poly in mask.polygons
]
return polygons, mask
def get_labeled_seg(predictions,
score_threshold,
vis,
assigned_colors=None,
paper_img=False):
boxes = predictions.pred_boxes
scores = predictions.scores
classes = predictions.pred_classes
chosen = (scores > score_threshold).nonzero()[0]
boxes = boxes[chosen]
scores = scores[chosen]
classes = classes[chosen]
# labels = list(range(len(predictions)))
labels = [f"{idx}: {score:.2f}" for idx, score in enumerate(scores)]
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, vis.output.height, vis.output.width) for x in masks
]
alpha = 0.5
if vis._instance_mode == ColorMode.IMAGE_BW:
vis.output.img = vis._create_grayscale_image(
(predictions.pred_masks.any(dim=0) > 0).numpy())
alpha = 0.3
if paper_img:
boxes = None
labels = None
vis.overlay_instances(
masks=masks,
assigned_colors=assigned_colors,
boxes=boxes,
labels=labels,
alpha=alpha,
)
seg_pred = vis.output.get_image()
return seg_pred
def draw_instance_predictions(self, predictions):
"""
:param predictions:
:return: Besides the functions of its mother class method, this method deals with extreme points.
"""
ext_points = predictions.ext_points if predictions.has(
"ext_points") else None
pred_polys = predictions.pred_polys if predictions.has(
"pred_polys") else None
if False:
return super().draw_instance_predictions(predictions)
else:
boxes = predictions.pred_boxes if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
labels = _create_text_labels(
classes, scores, self.metadata.get("thing_classes", None))
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
if predictions.has("pred_masks"):
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, self.output.height, self.output.width)
for x in masks
]
else:
if predictions.has("pred_polys"):
output_height = predictions.image_size[0]
output_width = predictions.image_size[1]
pred_masks = get_polygon_rles(
predictions.pred_polys.flatten(),
(output_height, output_width))
masks = np.asarray(pred_masks)
masks = [
GenericMask(x, self.output.height, self.output.width)
for x in masks
]
else:
masks = None
path = predictions.pred_path.numpy() if predictions.has(
"pred_path") else None
if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get(
"thing_colors"):
colors = [
self._jitter(
[x / 255 for x in self.metadata.thing_colors[c]])
for c in classes
]
alpha = 0.8
else:
colors = None
alpha = 0.5
if self._instance_mode == ColorMode.IMAGE_BW:
assert predictions.has(
"pred_masks"), "ColorMode.IMAGE_BW requires segmentations"
self.output.img = self._create_grayscale_image(
(predictions.pred_masks.any(dim=0) > 0).numpy())
alpha = 0.3
self.overlay_instances(
masks=masks,
boxes=boxes,
labels=labels,
ext_points=ext_points,
path=path,
keypoints=keypoints,
assigned_colors=colors,
alpha=alpha,
)
return self.output
def run_on_image(self, image, debug):
"""
Args:
image (np.ndarray): an image of shape (H, W, C) (in BGR order).
This is the format used by OpenCV.
Returns:
predictions (dict): the output of the model.
vis_output (VisImage): the visualized image output.
"""
vis_output = None
obj = None
predictions = self.predictor(image.astype(np.uint8))
# Convert image from OpenCV BGR format to Matplotlib RGB format.
image = image[:, :, ::-1]
visualizer = Visualizer(image,
self.metadata,
instance_mode=self.instance_mode)
if "panoptic_seg" in predictions:
panoptic_seg, segments_info = predictions["panoptic_seg"]
vis_output = visualizer.draw_panoptic_seg_predictions(
panoptic_seg.to(self.cpu_device), segments_info)
if debug:
print('in panoptic_seg')
else:
if "sem_seg" in predictions:
vis_output = visualizer.draw_sem_seg(
predictions["sem_seg"].argmax(dim=0).to(self.cpu_device))
if debug:
print("in sem_seg")
if "instances" in predictions:
instances = predictions["instances"].to(self.cpu_device)
if debug:
vis_output = visualizer.draw_instance_predictions(
predictions=instances)
print('in instances')
#if output is json, debug is false
if not debug:
boxes = instances.pred_boxes.tensor.numpy(
) if instances.has("pred_boxes") else None
scores = instances.scores if instances.has(
'scores') else None
classes = instances.pred_classes if instances.has(
"pred_classes") else None
labels = _create_text_labels(
classes, scores,
visualizer.metadata.get("thing_classes", None))
keypoints = instances.pred_keypoints if instances.has(
"pred_keypoints") else None
if instances.has("pred_masks"):
masks = np.asarray(instances.pred_masks)
masks = [
GenericMask(x, visualizer.output.height,
visualizer.output.width) for x in masks
]
else:
masks = None
obj = {}
for i, _ in enumerate(labels):
tmp = {}
split = labels[i].split()
tmp['class'] = split[0]
tmp['score'] = scores[i].item()
tmp['box'] = {}
tmp['box']['left-up'] = [
boxes[i][0].item(), boxes[i][1].item()
]
tmp['box']['right-down'] = [
boxes[i][2].item(), boxes[i][3].item()
]
tmp['polygons'] = {}
if masks is not None:
for idx, segment in enumerate(masks[i].polygons):
tmp['polygons'][idx] = segment.reshape(
-1, 2).tolist()
obj[i] = tmp
return predictions, vis_output, obj
def main(args):
# # register dataset
# register_coco_instances(name="cause_brain_2019_train", metadata={},
# json_file="./datasets/coco/annotations/instances_train2019.json",
# image_root="./datasets/coco/train2019")
#
# register_coco_instances(name="cause_brain_2019_val", metadata={},
# json_file="./datasets/coco/annotations/instances_valid2019.json",
# image_root="./datasets/coco/valid2019")
#
# register_coco_instances(name="cause_brain_2019_test", metadata={},
# json_file="./datasets/coco/annotations/instances_test2019.json",
# image_root="./datasets/coco/test2019")
cfg = setup(args, eval=True)
# os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
# evaluate trained model
if not args.masks_only:
os.makedirs(os.path.join(cfg.OUTPUT_DIR, "inference"), exist_ok=True)
os.makedirs(os.path.join(cfg.OUTPUT_DIR, "masks"), exist_ok=True)
# # COCO OUTPUT:
#
# # setup trainer and evaluate
# trainer = DefaultTrainer(cfg)
# evaluator = COCOEvaluator(dataset_name="balloon_val", cfg=cfg, distributed=False, output_dir="./output/")
# val_loader = build_detection_test_loader(cfg, "balloon_val")
# print(inference_on_dataset(trainer.model, val_loader, evaluator))
# # another equivalent way is to use trainer.test
# MASK OUTPUT:
predictor = DefaultPredictor(cfg)
# dataset_dicts = DatasetCatalog.get(cfg.DATASETS.TEST[0])
metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
class_names = metadata.get("thing_classes", None)
for filename in get_paths(args.input):
basename = os.path.basename(filename)
# logger.info("Loading: " + filename)
im = cv2.imread(filename)
outputs = predictor(im)
if len(outputs['instances']) > 0:
if len(outputs['instances']) > 0:
instances = outputs["instances"].to("cpu")
scores = instances.scores
logger.info("Scores" + str(scores))
classes = instances.pred_classes.numpy()
# labels = _create_text_labels(classes, scores, metadata.get("thing_classes", None))
labels = [class_names[i] for i in classes]
# logger.info("MASKS {}".format(instances.pred_masks))
try:
masks = np.asarray(instances.pred_masks)
masks = [
GenericMask(x, im.shape[0], im.shape[1]) for x in masks
]
# print(filename)
# print("Masks")
# print(labels)
for i in range(len(masks)):
name, ext = tuple(os.path.splitext(basename))
output_filename = name + "_" + labels[i] + ext
output_path = os.path.join(
cfg.OUTPUT_DIR, "" if args.masks_only else "masks",
output_filename)
draw_polygons(im.shape[0], im.shape[1],
masks[i].polygons, output_path)
# for segment in masks[i].polygons:
# print(i, segment.reshape(-1, 2).shape, labels[i])
if not args.masks_only:
vis = Visualizer(
im[:, :, ::-1],
metadata=metadata,
scale=2.0,
)
vis = vis.draw_instance_predictions(
outputs["instances"].to("cpu"))
vis.save(
os.path.join(cfg.OUTPUT_DIR, "inference",
basename))
except AttributeError:
logger.error("No segmentation exists!")
else:
logger.info('No Instances found!')
def main(dataRoot, testDir, annoRoot, taskName):
testFunc = createDataFunc(testDir, annoRoot)
DatasetCatalog.register("TestSetBatch", testFunc)
MetadataCatalog.get('TestSetBatch').set(
thing_classes=['Houses', 'Buildings', 'Sheds/Garages'])
visRoot = os.path.join(dataRoot, 'visRes')
predVisRoot = os.path.join(visRoot, 'predRes')
predPlVisRoot = os.path.join(visRoot, 'predPolyRes')
predMkVisRoot = os.path.join(visRoot, 'predMaskRes')
if not os.path.exists(visRoot):
os.mkdir(visRoot)
if not os.path.exists(predVisRoot):
os.mkdir(predVisRoot)
if not os.path.exists(predPlVisRoot):
os.mkdir(predPlVisRoot)
if not os.path.exists(predMkVisRoot):
os.mkdir(predMkVisRoot)
cfg = get_cfg()
cfg.merge_from_file(
"Cfg/COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml")
cfg.DATASETS.TRAIN = ("TrainSetBatch", )
cfg.DATALOADER.NUM_WORKERS = 24
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.OUTPUT_DIR = './' + taskName + '_output'
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 #测试的阈值
predictor = DefaultPredictor(cfg)
import cv2
import matplotlib.pyplot as plt
from detectron2.utils.visualizer import Visualizer, ColorMode, GenericMask
dataset_dicts = testFunc()
meta_dicts = MetadataCatalog.get('TestSetBatch')
from rdp import rdp
for ii in range(0, len(dataset_dicts)):
img = cv2.imread(dataset_dicts[ii]["file_name"])
vPred = Visualizer(img[:, :, ::-1],
metadata=meta_dicts,
scale=1.0,
instance_mode=ColorMode.IMAGE_BW)
vPoly = Visualizer(img[:, :, ::-1],
metadata=meta_dicts,
scale=1.0,
instance_mode=ColorMode.IMAGE_BW)
vMask = Visualizer(img[:, :, ::-1],
metadata=meta_dicts,
scale=1.0,
instance_mode=ColorMode.IMAGE_BW)
outputs = predictor(img)
res = outputs["instances"].to("cpu")
vPred.output = vPred.draw_instance_predictions(res)
picName = str.split(dataset_dicts[ii]["file_name"], '/')[-1]
picName = os.path.join(predVisRoot, picName)
plt.imsave(picName, vPred.output.get_image()[:, :, ::-1])
height = res.image_size[0]
width = res.image_size[1]
masks = [
GenericMask(x.numpy(), height, width).mask_to_polygons(x.numpy())
for x in res.pred_masks
]
masksBin = [
GenericMask(x.numpy(), height, width).mask for x in res.pred_masks
]
labels = [int(x.numpy()) for x in res.pred_classes]
count = 0
colorMaps = [[0., 0., 1.], [1., 0., 0.], [1.0, 0.3, .5]]
for mask in masks:
colorIn = colorMaps[labels[count]]
maskRaw = mask[0][0].reshape(-1, 2)
maskSim = rdp(maskRaw, epsilon=6)
vPoly.output = vPoly.draw_polygon(maskSim, colorIn, alpha=0.3)
vMask.output = vMask.draw_binary_mask(masksBin[count],
colorIn,
alpha=0.3)
count += 1
picName = str.split(dataset_dicts[ii]["file_name"], '/')[-1]
picName = os.path.join(predPlVisRoot, picName)
plt.imsave(picName, vPoly.output.get_image()[:, :, ::-1])
picName = str.split(dataset_dicts[ii]["file_name"], '/')[-1]
picName = os.path.join(predMkVisRoot, picName)
plt.imsave(picName, vMask.output.get_image()[:, :, ::-1])
def draw_instance_predictions(self, predictions, category=None):
"""
Draw instance-level prediction results on an image.
Args:
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
category: the integer category for the desired annotation to display as a list or None if all of them
Returns:
output (VisImage): image object with visualizations.
"""
# start additional code
if category == None:
boxes = predictions.pred_boxes if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
labels = self._create_text_labels(
classes, scores, self.metadata.get("thing_classes", None))
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
else:
all_boxes = predictions.pred_boxes if predictions.has(
"pred_boxes") else None
all_scores = predictions.scores if predictions.has(
"scores") else None
all_classes = predictions.pred_classes if predictions.has(
"pred_classes") else None
all_labels = self._create_text_labels(
all_classes, all_scores,
self.metadata.get("thing_classes", None))
all_keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
boxes = [] if all_boxes != None else None
scores = [] if all_scores != None else None
classes = [] if all_classes != None else None
labels = [] if all_labels != None else None
keypoints = [] if all_keypoints != None else None
for c in category:
for i in range(0, len(all_classes)):
if all_classes[i] == c:
classes.append(all_classes[i])
if all_boxes != None:
boxes.append(all_boxes[i])
if all_scores != None:
scores.append(all_scores[i])
if all_labels != None:
labels.append(all_labels[i])
if all_keypoints != None:
keypoints.append(all_keypoints[i])
if boxes != None and len(boxes) > 0:
boxes = Boxes(torch.cat([b.tensor for b in boxes], dim=0))
if scores != None and len(scores) > 0:
scores = torch.stack(scores)
if classes != None and len(classes) > 0:
classes = torch.stack(classes)
# end additional code
# removed alpha from here and put it as fixed value
if predictions.has("pred_masks"):
masks = np.asarray(predictions.pred_masks)
masks = [
GenericMask(x, self.output.height, self.output.width)
for x in masks
]
else:
masks = None
if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get(
"thing_colors"):
colors = [
self._jitter([x / 255 for x in self.metadata.thing_colors[c]])
for c in classes
]
else:
colors = None
if self._instance_mode == ColorMode.IMAGE_BW:
self.output.img = self._create_grayscale_image(
(predictions.pred_masks.any(dim=0) > 0).numpy())
self.overlay_instances(labels=labels,
boxes=boxes,
masks=masks,
keypoints=keypoints,
assigned_colors=colors,
alpha=1)
return self.output
def iou_metrics(det_masks: List[np.ndarray],
gt_masks: List[np.ndarray],
tp_threshold=0.5,
det_bboxes=None,
gt_bboxes=None):
"""
iou metric in one frame
:param dets:
:param gts:
:param tp_threshold:
:return:
"""
FNs = []
FPs = []
TPs = []
# create iou mapping. only above threshold objects are recorded
iou_map = {} # [i_det][i_gt]=iou
iou_rev_map = {} # [i_gt][i_det]=iou
# create bboxes if bboxes are not provided
if gt_bboxes is None:
gt_bboxes = [GenericMask(m, *m.shape).bbox() for m in gt_masks]
if det_bboxes is None:
det_bboxes = [GenericMask(m, *m.shape).bbox() for m in det_masks]
for i_gt, i_det in product(range(len(gt_masks)), range(len(det_masks))):
gt_mask = gt_masks[i_gt]
det_mask = det_masks[i_det]
gt_bbox = gt_bboxes[i_gt]
det_bbox = det_bboxes[i_det]
# use bbox to determine if iou should be calculated
if bbox_iou(gt_bbox, det_bbox) < tp_threshold:
continue
iou_value = compute_iou(gt_mask, det_mask)
if iou_value >= tp_threshold:
iou_map.setdefault(i_det, {})[i_gt] = iou_value
iou_rev_map.setdefault(i_gt, {})[i_det] = iou_value
for i_gt, _ in enumerate(gt_masks):
# False Negative: gt no det
if i_gt not in iou_rev_map:
FNs.append(i_gt)
else:
dets_of_gt: dict = iou_rev_map[i_gt]
if len(dets_of_gt) == 1:
# single TP
TPs.append((list(dets_of_gt.keys())[0], i_gt))
else:
# multiple det on same GT
dets_sorted = sorted(dets_of_gt.items(), key=lambda x: x[1])
TPs.append((dets_sorted[0][0], i_gt))
for det_lower_iou, _ in enumerate(dets_sorted[1:]):
FPs.append(det_lower_iou)
for i_det, _ in enumerate(det_masks):
# False Positive: det no gt
if i_det not in iou_map:
FPs.append(i_det)
tp_ious = [iou_map[tp[0]][tp[1]] for tp in TPs]
# precision:
precision_ious = [0] * len(FPs) + tp_ious
precision = np.mean(precision_ious)
# recall:
common_recall_ious = [0] * len(FNs) + [
1 if x > tp_threshold else 0 for x in tp_ious
]
common_recall = np.mean(common_recall_ious)
recall_ious = [0] * len(FNs) + tp_ious
recall = np.mean(recall_ious)
return {
"FPs": FPs,
"TPs": TPs,
"FNs": FNs,
"iou_map": iou_map,
"iou_rev_map": iou_rev_map,
"precision": precision,
"recall": recall,
"common_recall": common_recall,
"precision_ious": precision_ious,
"recall_ious": recall_ious,
"common_recall_ious": common_recall_ious,
}