Skip to content

Feature/erosion #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chinzhiweiblank wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Feature/erosion #5

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d981786
Add erosion and split checks
chinzhiweiblank Jan 30, 2020
328850d
Add checks for errant masks
chinzhiweiblank Jan 30, 2020
b67fe81
Add file name extensions
chinzhiweiblank Jan 30, 2020
21d34cb
Add conversion from eroded mask to polygon
chinzhiweiblank Feb 7, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
180 changes: 124 additions & 56 deletions visualize_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import matplotlib.pyplot as plt
from skimage import measure
import cv2
import numpy as np
import pandas as pd
Expand All @@ -8,27 +8,25 @@
import json
import copy
import argparse

import scipy.ndimage.measurements
from tensorpack.utils import logger, viz
from tensorpack.utils.timer import timed_operation
from tensorpack.utils.palette import PALETTE_RGB

from pycocotools import mask as maskUtils

import pycocotools.mask as cocomask
from six.moves import zip


class COCODetection(object):
# handle the weird (but standard) split of train and val

# Not used
_INSTANCE_TO_BASEDIR = {
'valminusminival2014': 'val2014',
'minival2014': 'val2014',
}

COCO_id_to_category_id = {1: 1, 2: 2, 3: 3, 5: 4, 6: 5}
category_id_to_COCO_id = {v:k for k,v in COCO_id_to_category_id.items()}
category_id_to_COCO_id = {v: k for k, v in COCO_id_to_category_id.items()}
"""
Mapping from the incontinuous COCO category id to an id in [1, #category]
For your own dataset, this should usually be an identity mapping.
Expand Down Expand Up @@ -60,7 +58,8 @@ def print_coco_metrics(self, json_file):
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
fields = ['IoU=0.5:0.95', 'IoU=0.5', 'IoU=0.75', 'small', 'medium', 'large']
fields = ['IoU=0.5:0.95', 'IoU=0.5',
'IoU=0.75', 'small', 'medium', 'large']
for k in range(6):
ret['mAP(bbox)/' + fields[k]] = cocoEval.stats[k]

Expand Down Expand Up @@ -90,7 +89,7 @@ def load(self, add_gt=True, add_mask=False):
assert add_gt
with timed_operation('Load Groundtruth Boxes for {}'.format(self.name)):
img_ids = self.coco.getImgIds()
img_ids.sort()
#img_ids.sort()
# list of dict, each has keys: height,width,id,file_name
imgs = self.coco.loadImgs(img_ids)

Expand All @@ -116,7 +115,8 @@ def _add_detection_gt(self, img, add_mask):
"""
# ann_ids = self.coco.getAnnIds(imgIds=img['image_id'])
# objs = self.coco.loadAnns(ann_ids)
objs = self.coco.imgToAnns[img['image_id']] # equivalent but faster than the above two lines
# equivalent but faster than the above two lines
objs = self.coco.imgToAnns[img['image_id']]

# clean-up boxes
valid_objs = []
Expand Down Expand Up @@ -145,20 +145,25 @@ def _add_detection_gt(self, img, add_mask):
assert obj['iscrowd'] == 1
obj['segmentation'] = None
else:
valid_segs = [np.asarray(p).reshape(-1, 2).astype('float32') for p in segs if len(p) >= 6]
valid_segs = [np.asarray(
p).reshape(-1, 2).astype('float32') for p in segs if len(p) >= 6]
if len(valid_segs) == 0:
logger.error("Object {} in image {} has no valid polygons!".format(objid, img['file_name']))
logger.error("Object {} in image {} has no valid polygons!".format(
objid, img['file_name']))
elif len(valid_segs) < len(segs):
logger.warn("Object {} in image {} has invalid polygons!".format(objid, img['file_name']))
logger.warn("Object {} in image {} has invalid polygons!".format(
objid, img['file_name']))

obj['segmentation'] = valid_segs

# all geometrically-valid boxes are returned
boxes = np.asarray([obj['bbox'] for obj in valid_objs], dtype='float32') # (n, 4)
boxes = np.asarray([obj['bbox']
for obj in valid_objs], dtype='float32') # (n, 4)
cls = np.asarray([
self.COCO_id_to_category_id[obj['category_id']]
for obj in valid_objs], dtype='int32') # (n,)
is_crowd = np.asarray([obj['iscrowd'] for obj in valid_objs], dtype='int8')
is_crowd = np.asarray([obj['iscrowd']
for obj in valid_objs], dtype='int8')

# add the keys
img['boxes'] = boxes # nx4
Expand All @@ -168,10 +173,10 @@ def _add_detection_gt(self, img, add_mask):
# also required to be float32
img['segmentation'] = [
obj['segmentation'] for obj in valid_objs]

def getClassNameFromSample(self, class_id):
return self.coco.loadCats(self.category_id_to_COCO_id[int(class_id)])[0]["name"]

@staticmethod
def load_many(basedir, names, add_gt=True, add_mask=False):
"""
Expand All @@ -186,35 +191,39 @@ def load_many(basedir, names, add_gt=True, add_mask=False):
coco = COCODetection(basedir, n)
ret.extend(coco.load(add_gt, add_mask=add_mask))
return ret



def getClassesFromImg(img):
return img["class"]


def getMasksFromImg(img):
is_crowd = img['is_crowd']
segmentation = copy.deepcopy(img['segmentation'])
segmentation = [segmentation[k] for k in range(len(segmentation)) if not is_crowd[k]]
segmentation = [segmentation[k]
for k in range(len(segmentation)) if not is_crowd[k]]
height, width = img['height'], img['width']
# Apply augmentation on polygon coordinates.
# And produce one image-sized binary mask per box.
masks = []
width_height = np.asarray([width, height], dtype=np.float32)
for polys in segmentation:
# if not cfg.DATA.ABSOLUTE_COORD:
# polys = [p * width_height for p in polys]
# polys = [aug.augment_coords(p, params) for p in polys]
masks.append(segmentation_to_mask(polys, height, width))
masks.append(segmentation_to_mask(polys, height, width,
linear=(img['category_ids'] == [1])))
masks = np.asarray(masks, dtype='uint8') # values in {0, 1}
return masks


def genBoxesFromMasks(masks):
"""Compute bounding boxes from masks.
mask: [num_instances, height, width]. Mask pixels are either 1 or 0.
Returns: bbox array [num_instances, (y1, x1, y2, x2)].
"""
boxes = np.zeros([masks.shape[0], 4], dtype=np.int32)
for i in range(masks.shape[0]):
m = masks[i ,:, :]
m = masks[i, :, :]
# Bounding box.
horizontal_indicies = np.where(np.any(m, axis=0))[0]
vertical_indicies = np.where(np.any(m, axis=1))[0]
Expand All @@ -232,23 +241,38 @@ def genBoxesFromMasks(masks):
return boxes.astype(np.int32)


def segmentation_to_mask(polys, height, width):
def segmentation_to_mask(polys, height, width, linear=False):
"""
Convert polygons to binary masks.
Args:
polys: a list of nx2 float array. Each array contains many (x, y) coordinates.
height, width: dimensions of segmentation
linear: Boolean for erosion of linear cracks
Returns:
a binary matrix of (height, width)
"""
polys = [p.flatten().tolist() for p in polys]
assert len(polys) > 0, "Polygons are empty!"

import pycocotools.mask as cocomask
rles = cocomask.frPyObjects(polys, height, width)
rle = cocomask.merge(rles)
return cocomask.decode(rle)
res_rle = cocomask.decode(rle)
#print(mask_to_polygons(res_rle))
if linear:
return cv2.erode(res_rle, np.ones((6, 6), np.uint8))
else:
return res_rle

def mask_to_polygons(mask):
polygons = []
contours = measure.find_contours(mask, 0.5)
for contour in contours:
contour = np.flip(contour, axis=1)
segmentation = contour.ravel().tolist()
polygons.append(segmentation)
return polygons

def draw_mask(im, mask, box, label, alpha=0.5, color=None):
def draw_mask(im, mask, box, label, alpha=0.5, color=None, linear=False):
"""
Overlay a mask on top of the image.
Args:
Expand All @@ -263,41 +287,85 @@ def draw_mask(im, mask, box, label, alpha=0.5, color=None):
im = im.astype('uint8')
color_tuple = tuple([int(c) for c in color])
im = viz.draw_boxes(im, box[np.newaxis, :], [label], color=color_tuple)
return im

cc = 1
if linear:
_, cc = scipy.ndimage.measurements.label(
mask, structure=np.ones((3, 3)))
return cc == 1, im


def parse_args():
parser = argparse.ArgumentParser(description='Code for Harris corner detector tutorial.')
parser = argparse.ArgumentParser(
description='Code for Harris corner detector tutorial.')
parser.add_argument('--imagedir', help='Path to dataset images.')
parser.add_argument('--jsonfile', help='Path to json file.')
parser.add_argument('--output')
parser.add_argument('--check', help='Flag to purely check JSON', action='store_true', default = False)
parser.add_argument('--output', help='Output Directory for images with masks', default='output_dir')
parser.add_argument('--outputjson', help='Output JSON', default="output.json")
return parser.parse_args()



def main():
args = parse_args()
output_dir = args.output
ds = COCODetection(args.imagedir,args.jsonfile)
imgs = ds.load(add_gt=True, add_mask=True)
os.makedirs(output_dir, exist_ok=True)
for img in tqdm.tqdm(imgs):
# Get masks from "img" (it's actually the image's meta rather than the image itself)
# I follow the same naming from the Tensorpack's implementation of COCODetection
masks = getMasksFromImg(img)
boxes = genBoxesFromMasks(masks)
classes = getClassesFromImg(img) # Class IDs
classes = [ds.getClassNameFromSample(clsId) for clsId in classes] # Class names
file_name = img['file_name']
image_id = img['image_id']
im = cv2.imread(file_name)
orig_im = im.copy()
# Draw masks, boxes and labels
for i in range(masks.shape[0]):
im = draw_mask(im, masks[i], boxes[i], str(classes[i]))
basename = os.path.basename(file_name)

output_path = os.path.join(output_dir, str(image_id) + '_' + basename)
# merge original image to the image with labels
im = np.concatenate([orig_im, im], axis=1)
cv2.imwrite(output_path, im)
errant_imgs = set()
args = parse_args()
if args.check:
ds = COCODetection(args.imagedir, args.jsonfile)
imgs = ds.load(add_gt=True, add_mask=True)
for img in tqdm.tqdm(imgs):
masks = getMasksFromImg(img)
image_id = img['image_id']
basename = os.path.basename(img['file_name'])
for mask in masks:
if 1 in img['category_ids'] or 2 in img['category_ids']:
_, cc = scipy.ndimage.measurements.label(mask, structure=np.ones((3, 3)))
if cc!=1: errant_imgs.add(str(image_id)+'_'+basename)
else:
annotations = []
output_dir = args.output
ds = COCODetection(args.imagedir, args.jsonfile)
imgs = ds.load(add_gt=True, add_mask=True)
with open(args.jsonfile, 'r') as jsonfile: data = json.load(jsonfile)
os.makedirs(output_dir, exist_ok=True)
for img in tqdm.tqdm(imgs):
# Get masks from "img" (it's actually the image's meta rather than the image itself)
# I follow the same naming from the Tensorpack's implementation of COCODetection
masks = getMasksFromImg(img)
boxes = genBoxesFromMasks(masks)
classes = getClassesFromImg(img) # Class IDs
classes = [ds.getClassNameFromSample(
clsId) for clsId in classes] # Class names
file_name = img['file_name']
image_id = img['image_id']
im = cv2.imread(file_name)
orig_im = im.copy()
# Draw masks, boxes and labels
# For images with a linear crack, erosion is performed
for i in range(masks.shape[0]):
connected, im = draw_mask(im, masks[i], boxes[i], str(
classes[i]), linear=(img['category_ids'] == [1]))
annotations.append(mask_to_polygons(masks[i]))
if connected == False:
errant_imgs.add(img['path'])

basename = os.path.basename(file_name)
output_path = os.path.join(output_dir, str(image_id) + '_' + basename)

# merge original image to the image with labels
im = np.concatenate([orig_im, im], axis=1)
cv2.imwrite(output_path, im)
print(len(annotations))
with open(args.jsonfile, 'r') as json_file: data = json.load(json_file)
for i in range(len(annotations)):
data["annotations"][i]['segmentation']=annotations[i]
with open(args.output_json, 'w') as json_file: json.dump(data, json_file)
# Errant Images where mask erosion separated the cracks
if len(errant_imgs)!=0:
print(f"Number of Errant Images: {len(errant_imgs)}")
print("List of Errant Images:")
for img_err in errant_imgs:
print(img_err)
else:
print("No errant images")

if __name__ == '__main__':
main()