def test(cfg):
predictor = DefaultPredictor(cfg)
ret = []
df = json.load(open('dataset/test.json'))
for image in df['images']:
file_name = image['file_name']
height = image['height']
width = image['width']
print("testing ", file_name)
file_name = os.path.join("dataset/test_images", file_name)
image_id = image['id']
img = cv2.imread(file_name)
output = predictor(img)
info = output['instances'].get_fields()
num_inst = len(info['scores'].to("cpu"))
for i in range(num_inst):
instance = {}
instance['image_id'] = image_id
instance['score'] = float(info['scores'][i].to("cpu").numpy())
instance['category_id'] = int(
info['pred_classes'][i].to("cpu").numpy() + 1)
bit_mask = info['pred_masks'][i].to("cpu").numpy()
rle = binary_mask_to_rle(bit_mask)
rle['size'] = [height, width]
instance['segmentation'] = rle
ret.append(instance)
return ret
def create_submission(json, dir, model):
cocoGt = COCO(json)
coco_dt = []
with torch.no_grad():
model.eval()
for imgid in cocoGt.imgs:
print(imgid)
img_path = os.path.join(dir,
cocoGt.loadImgs(ids=imgid)[0]['file_name'])
image = Image.open(img_path).convert("RGB")
image = torchvision.transforms.functional.to_tensor(image.copy())
image = image.to(device)
prediction = model([image])[0]
masks = prediction['masks'].cpu()
masks.squeeze_(1)
categories = prediction['labels'].cpu()
scores = prediction['scores'].cpu()
n_instances = len(scores)
if len(categories
) > 0: # If any objects are detected in this image
for i in range(n_instances): # Loop all instances
# save information of the instance in a dictionary then append on coco_dt list
pred = {}
pred[
'image_id'] = imgid # this imgid must be same as the key of test.json
pred['category_id'] = int(categories[i])
mask = masks[i, :, :].numpy()
mask[mask > 0.5] = 1
mask[mask <= 0.5] = 0
pred['segmentation'] = binary_mask_to_rle(
mask) # save binary mask to RLE, e.g. 512x512 -> rle
pred['score'] = float(scores[i])
coco_dt.append(pred)
return coco_dt
def build_coco_results(dataset, image_ids, rois, class_ids, scores, masks):
"""Arrange resutls to match COCO specs in http://cocodataset.org/#format
"""
# If no results, return an empty list
if rois is None:
return []
results = []
for image_id in image_ids:
# Loop through detections
for i in range(rois.shape[0]):
class_id = class_ids[i]
score = scores[i]
bbox = np.around(rois[i], 1)
mask = masks[:, :, i]
pred = {}
# this imgid must be same as the key of test.json
pred['image_id'] = image_id
pred['category_id'] = dataset.get_source_class_id(class_id, "coco")
# save binary mask to RLE, e.g. 512x512 -> rle
pred['segmentation'] = binary_mask_to_rle(mask)
pred['score'] = float(score)
results.append(pred)
# result = {
# "image_id": image_id,
# "category_id": dataset.get_source_class_id(class_id, "coco"),
# #"bbox": [bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0]],
# "score": score,
# #"segmentation": maskUtils.encode(np.asfortranarray(mask))
# "segmentation": binary_mask_to_rle(masks[:,:,i])
# }
# results.append(result)
return results
def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Batch_size = 8
# '''
version = args.version
print(version)
cfg = get_cfg_defaults()
cfg.merge_from_file(args.config)
cfg.freeze()
trainable = cfg.TRAIN.TRAINABLE
class_number = cfg.TRAIN.CLASS_NUMBER
hidden_layer = cfg.TRAIN.HIDDEN_LAYER
isnorm = cfg.TRAIN.ISNORM
model = function.get_instance_segmentation_model(class_number, trainable,
hidden_layer)
load_model_path = os.path.join(args.model_path, version, 'model',
version + str(args.epoch))
model.load_state_dict(
torch.load(load_model_path, map_location=lambda storage, loc: storage))
model.to(device)
# '''
print("test dataset!")
if isnorm:
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
else:
test_transform = transforms.Compose([
transforms.ToTensor(),
])
test_dataset = load_data.TPVDataset_test(test_transform, args.test)
test_loader = Data.DataLoader(test_dataset,
batch_size=8,
shuffle=False,
num_workers=8,
collate_fn=utils.collate_fn)
prediction = []
with torch.no_grad():
model.eval()
for batch, (images, idx) in enumerate(test_loader, 1):
images = list(img.to(device) for img in images)
outputs = model(images)
n_instances = len(outputs)
for i, p in enumerate(outputs):
for label, mask, score in zip(p['labels'], p['masks'],
p['scores']):
pred = {}
pred['image_id'] = idx[i]
pred['category_id'] = int(label.cpu())
np.set_printoptions(threshold=np.inf)
mask = np.where(mask.cpu().numpy() > args.mask_threshold,
1, 0).astype(np.uint8)[0, :, :]
pred['segmentation'] = utils.binary_mask_to_rle(mask)
pred['score'] = float(score.cpu())
prediction.append(pred)
prediction_path = './prediction'
if not os.path.isdir(prediction_path):
os.makedirs(prediction_path)
import json
with open('./prediction/' + '0856165_' + args.output_version + '.json',
"w") as f:
json.dump(prediction, f)
with open('./prediction/' + 'match.txt', "a") as f:
print(args.output_version,
args.version,
args.epoch,
" -",
args.mask_threshold,
file=f)
print("END")
for imgid in range(100):
image = read_image_bgr("test_images/" +
coco.loadImgs(ids=imgid)[0]['file_name'])
image = preprocess_image(image)
image, scale = resize_image(image)
outputs = model.predict_on_batch(np.expand_dims(image, axis=0))
boxes = outputs[-4][0]
scores = outputs[-3][0]
labels = outputs[-2][0]
masks = outputs[-1][0]
# correct for image scale
boxes /= scale
# visualize detections
for box, score, label, mask in zip(boxes, scores, labels, masks):
if score < 0.5:
break
pred = {}
pred[
'image_id'] = imgid # this imgid must be same as the key of test.json
pred['category_id'] = label
pred['segmentation'] = binary_mask_to_rle(
mask) # save binary mask to RLE, e.g. 512x512 -> rle
pred['score'] = score
coco_dt.append(pred)
with open("submission.json", "w") as f:
json.dump(coco_dt, f)
labels = outputs[0]["labels"]
boxes = outputs[0]["boxes"]
masks = np.round(masks.cpu())
n_instances = len(scores)
if (len(labels) >
0): # If any objects are detected in this image
for i in range(n_instances): # Loop all instances
# save information of the instance in a dictionary then append on coco_dt list
pred = {}
pred[
"image_id"] = imgid # this imgid must be same as the key of test.json
pred["score"] = float(scores[i])
pred["category_id"] = int(labels[i])
mask = masks[i][0].unsqueeze(2)
pred["segmentation"] = binary_mask_to_rle(
(masks[i][0]).detach().cpu().numpy(
)) # save binary mask to RLE, e.g. 512x512 -> rle
coco_dt.append(pred)
with open(cfg.result_pth + cfg.train_result, "w") as f:
json.dump(coco_dt, f)
cocoGt = COCO(cfg.coco_train_file)
cocoDt = cocoGt.loadRes(cfg.result_pth + cfg.train_result)
cocoEval = COCOeval(cocoGt, cocoDt, "segm")
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
print("Done!")
def test(cfg_path, pretrain_imageNet=True):
# info
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
test_info = parse_info_cfg(cfg_path)
print(test_info)
os.makedirs(test_info['img_path'], exist_ok=True)
# load data
test_data = dataLoader(test_info['src_path'], test_info['mode'],
test_info['img_size'])
test_loader = Data.DataLoader(dataset=test_data,
batch_size=test_info['batch'],
shuffle=False,
num_workers=test_info['n_cpu'],
collate_fn=test_data.collate_fn)
# load an instance segmentation model pre-trained on ImageNet
model = torchvision.models.detection.maskrcnn_resnet50_fpn(
num_classes=test_info['n_class'],
pretrained_backbone=pretrain_imageNet).to(device)
# test
submit_seq = []
model.load_state_dict(
torch.load(test_info['pkl_path'], map_location=device))
model.eval()
with torch.no_grad():
for imgs, _, img_ids, paths in tqdm.tqdm(test_loader):
imgs = imgs.to(device)
detections = model(imgs)
for detect_i, img_id_i, path in zip(detections, img_ids, paths):
detect_i = {
key: values.cpu()
for key, values in detect_i.items()
}
img_np = np.array(Image.open(path).convert('RGB'))
for j, l_j in enumerate(detect_i['labels']):
if detect_i['scores'][j] >= test_info['thres']:
detect_dict = {}
detect_dict['image_id'] = int(img_id_i)
detect_dict['score'] = float(detect_i['scores'][j])
detect_dict['category_id'] = int(l_j) + 1
_, out_mask = rescale(img_np.shape[1],
img_np.shape[0],
test_info['img_size'],
origin_mask=detect_i['masks'][j])
out_mask = torch.round(out_mask) # binary mask
detect_dict['segmentation'] = binary_mask_to_rle(
out_mask.numpy())
submit_seq.append(detect_dict)
if test_info['save_img']:
plot_img(
path, detect_i, test_info['img_size'],
f'{test_info["img_path"]}plot_{path.split("/")[-1]}',
test_info['class_names'])
# dump json file
json_path = test_info['saved_path'] + test_info['pkl_path'].split(
'/')[3].split('.')[
0] + f'_{test_info["mode"]}_{int(test_info["thres"]*100)}.json'
print(json_path)
with open(json_path, 'w') as outfile:
json.dump(submit_seq, outfile)
image = cv2.imread("test_images/" + cocoGt.loadImgs(
ids=imgid)[0]['file_name'])[:, :, ::-1] # load image
outputs = predictor(image)
# We can use `Visualizer` to draw the predictions on the image.
v = Visualizer(image[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imshow("myimage", v.get_image()[:, :, ::-1])
cv2.imwrite('save_image/' + cocoGt.loadImgs(ids=imgid)[0]['file_name'],
v.get_image()[:, :, ::-1])
# cv2.waitKey(0)
n_instances = len(outputs["instances"].scores)
if len(outputs["instances"].pred_classes
) > 0: # If any objects are detected in this image
for i in range(n_instances): # Loop all instances
# save information of the instance in a dictionary then append on coco_dt list
pred = {}
pred[
'image_id'] = imgid # this imgid must be same as the key of test.json
pred['category_id'] = int(outputs["instances"].pred_classes[i]) + 1
outputs["instances"].pred_masks[:, :, i].permute(0, 1)
pred['segmentation'] = binary_mask_to_rle(
outputs["instances"].pred_masks[i].to("cpu").numpy())
# save binary mask to RLE, e.g. 512x512 -> rle
pred['score'] = float(outputs["instances"].scores[i])
coco_dt.append(pred)
with open("0856052.json", "w") as f:
json.dump(coco_dt, f)
os.path.join(mask_dir, mask_filenames[d].text), 0)
polygons, binary_image = mask2poly((mask))
# Mask should be defined with at least 60 points.
area_binary_img = len(np.where(binary_image == 1)[0])
#if area_binary_img < area_thr:
#continue
# Write the new modified masks to a folder
cv2.imwrite(
os.path.join(dest_name, 'processed_masks',
mask_filenames[d].text),
binary_image * 255)
# Generate the RLE version of mask
RLE_mask = binary_mask_to_rle(mask)
RLE_mask_coco = comask.encode(
np.asfortranarray(binary_image.astype(np.uint8)))
# Area
area = comask.area(RLE_mask_coco)
# Iscrowd
if (len(mask_filenames) - sum_deleted) > 1:
iscrowd = 1
json_structure['annotations'].append({
'iscrowd':
iscrowd,
'bbox':
bbox_list,
'id':
cfg.MODEL.WEIGHTS = os.path.join(
cfg.OUTPUT_DIR, 'model_final.pth') # path to the model we just trained
predictor = DefaultPredictor(cfg)
coco_dt = []
for imgid in cocoGt.imgs:
image = cv2.imread(
os.path.join('./data/test_images',
cocoGt.loadImgs(ids=imgid)[0]['file_name']))
pred = predictor(image)['instances']
# masks, categories, scores = predictor(image) # run inference of your model
scores = pred.get('scores').to('cpu').numpy()
categories = pred.get('pred_classes').to('cpu').numpy()
masks = pred.get('pred_masks').to('cpu').numpy()
n_instances = len(scores)
if len(categories) > 0: # If any objects are detected in this image
for i in range(n_instances): # Loop all instances
# save information of the instance in a dictionary then append on coco_dt list
pred = {
'image_id': imgid,
'category_id': int(categories[i]) + 1,
'segmentation': binary_mask_to_rle(masks[i, :, :]),
'score': float(scores[i])
}
coco_dt.append(pred)
with open('submission.json', 'w') as f:
json.dump(coco_dt, f)
model = modellib.MaskRCNN(mode="inference", config=config, model_dir=MODEL_DIR)
model.load_weights(model.find_last(), by_name=True)
cocoGt = COCO("test.json")
coco_dt = []
"""# debug: output the image with masks (set imgid)
imgid = 34
image = cv2.imread("test_images/" + cocoGt.loadImgs(ids=imgid)[0]['file_name'])[:,:,::-1]
results = model.detect([image],verbose=1)
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
dataset_train.class_names, r['scores'])
plt.savefig(str(imgid)+'.png')
"""
for imgid in cocoGt.imgs:
image = cv2.imread("test_images/" +
cocoGt.loadImgs(ids=imgid)[0]['file_name'])[:, :, ::-1]
results = model.detect([image])
r = results[0]
n_instances = len(r['scores'])
for i in range(n_instances):
pred = {}
pred['image_id'] = imgid
pred['category_id'] = int(r['class_ids'][i])
pred['segmentation'] = binary_mask_to_rle(r['masks'][:, :, i])
pred['score'] = float(r['scores'][i])
coco_dt.append(pred)
with open('309551082.json', 'w') as f:
json.dump(coco_dt, f)
cfg.DATASETS.TEST = ("my_dataset")
predictor = DefaultPredictor(cfg)
cocoGt = COCO("test.json")
from utils import binary_mask_to_rle
coco_dt = []
for imgid in cocoGt.imgs:
image = cv2.imread("test_images/" + cocoGt.loadImgs(
ids=imgid)[0]['file_name'])[:, :, ::-1] # load image
outputs = predictor(image)
n_instances = len(outputs["instances"])
#print(n_instances)
if n_instances > 0: # If any objects are detected in this image
for i in range(n_instances): # Loop all instances
# save information of the instance in a dictionary then append on coco_dt list
pred = {}
pred[
'image_id'] = imgid # this imgid must be same as the key of test.json
pred['category_id'] = int(
outputs["instances"]._fields['pred_classes'][i]) + 1
#pred['segmentation'] = binary_mask_to_rle(masks[:,:,i]) # save binary mask to RLE, e.g. 512x512 -> rle
pred['segmentation'] = binary_mask_to_rle(outputs["instances"].to(
'cpu')._fields['pred_masks'][i].numpy())
pred['score'] = float(outputs["instances"]._fields['scores'][i])
coco_dt.append(pred)
with open("submission.json", "w") as f:
json.dump(coco_dt, f)