def main():
parser = init()
args = parser.parse_args()
# update the config options with the config file
cfg.merge_from_file(args.config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
# cfg.merge_from_list(["MODEL.MASK_ON", False])
coco_demo = COCODemo(
cfg,
confidence_threshold=args.confidence_threshold,
show_mask_heatmaps=args.show_mask_heatmaps,
masks_per_dim=args.masks_per_dim,
min_image_size=args.min_image_size,
)
if args.weight is not None:
checkpoint = torch.load(args.weight)
load_state_dict(coco_demo.model, checkpoint.pop("model"))
del checkpoint
start_time = time.time()
img = cv2.imread(args.img)
composite = coco_demo.run_on_opencv_image(img)
print("Time: {:.2f} s / img".format(time.time() - start_time))
# plt.figure()
plt.imshow(composite)
plt.savefig("result.jpg")
def viz(args):
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(["TEST.IMS_PER_BATCH", "1", "MODEL.WEIGHT", args.weight_file])
coco_demo = COCODemo(
cfg,
min_image_size=cfg.INPUT.MIN_SIZE_TEST,
confidence_threshold=float(args.confidence_threshold),
)
VAL_DATA_DIR = '../datasets/coco/timeline/val'
ANN_FILE = '../datasets/coco/timeline/val/coco_format.json'
OUTPUT_DIR = '../output/' + args.config_file.split('/')[-1].split('.')[0] + '-' + args.weight_file.split('/')[-1].split('.')[0]
if not os.path.exists(OUTPUT_DIR):
os.mkdir(OUTPUT_DIR)
gtDataset = coco.COCODataset(ANN_FILE, VAL_DATA_DIR, True)
for image, target, idx in gtDataset:
image = np.array(image)
top_prediction = coco_demo.compute_prediction(image)
top_prediction = coco_demo.select_top_predictions(top_prediction)
masked_image = draw_on_image(image, top_prediction)
target.add_field('mask', get_mask_array(target))
gt_image = draw_on_image(image, target)
cat_img = np.concatenate([image, masked_image, gt_image], axis=1)
imsave(cat_img, os.path.join(OUTPUT_DIR, '{}.png'.format(idx)))
print('finish {}'.format(idx))
def processing(threshold, jsonFile):
coco_demo = COCODemo(cfg,
min_image_size=800,
confidence_threshold=threshold)
imageList = []
for i, imagePath in enumerate(imagePathList):
print("loading:{}".format(i))
imagefilepath = os.path.join(IMAGE_PATH, imagePath)
image = Image.open(imagefilepath).convert("RGB")
image = np.array(image)[:, :, [2, 1, 0]]
imageList.append(image)
# compute predictions
predictions, errorList, successList = coco_demo.run_on_imageList(
imageList, imagePathList, jsonFile)
for i, prediction in enumerate(predictions):
print("writting:{}".format(i))
targetPath = os.path.join(TARGET_PATH, successList[i])
cv2.imwrite(targetPath, prediction)
imagePathList.remove(successList[i])
if successList[i] in errorPathList:
errorImagePath = os.path.join(ERROR_PATH, successList[i])
errorPathList.remove(successList[i])
os.remove(errorImagePath)
for i, name in enumerate(errorList):
print("writting:{}".format(i))
errorImagePath = os.path.join(IMAGE_PATH, name)
errorImage = cv2.imread(errorImagePath)
targetPath = os.path.join(ERROR_PATH, name)
cv2.imwrite(targetPath, errorImage)
def __init__(
self,
config_path=os.path.dirname(__file__) +
"/configs/caffe2/e2e_mask_rcnn_X_101_32x8d_FPN_1x_caffe2.yaml"):
self.model_config_path = config_path
cfg.merge_from_file(self.model_config_path)
cfg.merge_from_list([])
cfg.freeze()
self._greyscale_palette = (2 * 25 - 1)
# prepare object that handles inference plus adds predictions on top of image
self.coco_demo = COCODemo(cfg,
confidence_threshold=0.8,
show_mask_heatmaps=False,
masks_per_dim=1,
min_image_size=800)
self._greyscale_colours = self._generate_grayscale_values()
self.mask_rcnn_service = rospy.Service("MaskRcnnService", MaskRcnn,
self.mask_rcnn_service_callback)
print("Ready to receive mask rcnn service calls.")
self.mask_rcnn_publisher = rospy.Publisher('mask_rcnn_img',
Image,
queue_size=10)
def __init__(self, confidence_threshold=0.7):
cfg.merge_from_file('e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml')
cfg.MODEL.DEVICE
cfg.freeze()
self.model_wrapper = COCODemo(
cfg,
confidence_threshold=confidence_threshold,
)
def prepare_detector(self):
# Load model
cfg.merge_from_file(self.config)
cfg.merge_from_list(list())
cfg.MODEL.WEIGHT = self.weights
cfg.freeze()
self.detector = COCODemo(
cfg,
confidence_thresholds_for_classes=self.THRESHOLDS_FOR_CLASSES,
min_image_size=800)
def __init__(self, device):
self.config_file = "/maskrcnn-benchmark/configs/caffe2/e2e_faster_rcnn_R_101_FPN_1x_caffe2.yaml"
self.device = device
self.cfg = cfg
cfg.merge_from_file(self.config_file)
cfg.merge_from_list(["MODEL.DEVICE", device])
cfg.merge_from_list(["MODEL.MASK_ON", False])
self.coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.5,
)
def main():
parser = make_parser()
args = parser.parse_args()
# loading the mscoco validation set .
dataType = args.dataType
dataDir = args.dataDir
annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataType)
coco = COCO(annFile)
cats = coco.loadCats(coco.getCatIds())
#nms=[cat['name'] for cat in cats]
imgDir = '{}/images/{}'.format(dataDir, dataType)
#print('COCO categories: \n{}\n'.format(' '.join(nms)))
#nms = set([cat['supercategory'] for cat in cats])
#print('COCO supercategories: \n{}'.format(' '.join(nms)))
coco_demo = COCODemo(cfg, min_image_size=800, confidence_threshold=0.5)
coco_dataset = COCODataset(annFile, dataDir + '/images/' + args.dataType,
True)
initial_detection(coco_dataset,
coco_demo,
args,
max_images=args.max_images,
savedir=args.save_dir)
print(args)
#masked_detection( coco_dataset, coco_demo, args)
masked_detection_2(coco_dataset, coco_demo, args)
print('all_at_once:', args.all_at_once)
class FeatureExtractor():
def __init__(self, device):
self.config_file = "/maskrcnn-benchmark/configs/caffe2/e2e_faster_rcnn_R_101_FPN_1x_caffe2.yaml"
self.device = device
self.cfg = cfg
cfg.merge_from_file(self.config_file)
cfg.merge_from_list(["MODEL.DEVICE", device])
cfg.merge_from_list(["MODEL.MASK_ON", False])
self.coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.5,
)
def extract_features(self, batch_images, k=19):
box_list = self.coco_demo.extract_encoding_features(batch_images)
features = []
bboxes = []
for boxes in box_list:
encodings = boxes.get_field("encoding_features")
# Pad features with 0 to make them size k+1
encodings = np.pad(encodings, [(0, k+1-len(encodings)), (0, 0)], mode='constant')
# Pad boxes with 0 to make them size k
bbox = np.pad(boxes.bbox, [(0, k-len(boxes)), (0, 0)], mode='constant')
features.append(encodings)
bboxes.append(bbox)
return np.array(features), np.array(bboxes, dtype=np.int64)
def main():
config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda:0"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.5,
)
from pycocotools.coco import COCO
dataType = 'val2017'
dataDir = '/home/amir/data/coco'
annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataType)
coco = COCO(annFile)
cats = coco.loadCats(coco.getCatIds())
nms = [cat['name'] for cat in cats]
imgDir = '{}/images/{}'.format(dataDir, dataType)
print('COCO categories: \n{}\n'.format(' '.join(nms)))
nms = set([cat['supercategory'] for cat in cats])
print('COCO supercategories: \n{}'.format(' '.join(nms)))
class MaskRCNN(object):
def __init__(self, confidence_threshold=0.7):
cfg.merge_from_file('e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml')
cfg.MODEL.DEVICE
cfg.freeze()
self.model_wrapper = COCODemo(
cfg,
confidence_threshold=confidence_threshold,
)
def get_chips_and_masks(self, img, label_index=COCO_PERSON_INDEX):
'''
Params
------
img : nd array like, RGB
label_index : int, index of label wanted
Returns
-------
list of tuple (chip, mask)
- chip is a ndarray: bb crop of the image
- mask is a ndarray: same shape as chip, whose 'pixel' value is either 0 or 1, indicating if that pixel belongs to that class or not.
'''
preds = self.model_wrapper.compute_prediction(img)
top_preds = self.model_wrapper.select_top_predictions(preds)
labels = top_preds.get_field('labels')
person_bool_mask = (labels==label_index).numpy().astype(bool)
masks = top_preds.get_field('mask').numpy()[person_bool_mask]
bboxes = top_preds.bbox.to(torch.int64).numpy()[person_bool_mask]
results = []
for mask, box in zip( masks, bboxes ):
thresh = mask[0, :, :, None]
# l,t,r,b = box.to(torch.int64).numpy()
l,t,r,b = box
if b - t <= 0 or r - l <= 0:
continue
content = img[ t:(b+1), l:(r+1), : ]
minimask = thresh[ t:(b+1), l:(r+1), : ]
results.append( (content, minimask) )
return results
def detect_person(cfg, image):
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
)
predictions = coco_demo.compute_prediction(image)
top_predictions = coco_demo.select_top_predictions(predictions)
#result = coco_demo.overlay_class_names(result, top_predictions)
labels = top_predictions.get_field("labels").tolist()
labels = [coco_demo.CATEGORIES[i] for i in labels]
if 'person' in labels:
return 1
else:
return 0
def setup(opts):
config_file = "configs/caffe2/e2e_mask_rcnn_%s_1x_caffe2.yaml" % opts[
'architecture'].replace('-', '_')
cfg.merge_from_file(config_file)
if not torch.cuda.is_available():
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
model = COCODemo(
cfg,
confidence_threshold=opts['confidenceThreshold'],
)
return model
def setup():
config_file = "configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
cfg.merge_from_file(config_file)
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
)
return coco_demo
def main():
args = parser.parse_args()
cfg.merge_from_file(args.config)
cfg.merge_from_list(["MODEL.DEVICE", "cuda:0"])
coco_demo = COCODemo(
cfg,
min_image_size=2048,
confidence_threshold=0.7,
)
image = load_from_file(args.data)
batch_id = 1
elapsed = 0.
totaltime = 0.
while True:
try:
start_time = time.time()
predictions = coco_demo.run_on_opencv_image(image)
elapsed += time.time() - start_time
time.sleep(args.interval)
totaltime += time.time() - start_time
if batch_id % args.freq == 0:
print(
"Batch #{}: inference time: {:>3.3f}s, per image: {:>3.4f}s, total time: {:>3.3f}s inference percent: {:>3.0f}%"
.format(batch_id, elapsed, elapsed / args.freq, totaltime,
100. * elapsed / totaltime))
if args.output:
if not os.path.exists(args.output):
os.makedirs(args.output)
filename = os.path.join(
args.output, 'prediction_{}.jpg'.format(batch_id))
print("=> Saving prediction to {}".format(filename))
cv2.imwrite(filename, predictions)
elapsed = 0.
totaltime = 0.
batch_id += 1
except KeyboardInterrupt:
print("Interrupted.")
break
def single_predict():
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
config_file = "configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
#cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.5,
)
import cv2
from process_image import show_image, draw_bb
im = cv2.imread('/home/jianfw/data/sample_images/TaylorSwift.jpg',
cv2.IMREAD_COLOR)
predictions = coco_demo.compute_prediction(im)
predictions = coco_demo.select_top_predictions(predictions)
scores = predictions.get_field("scores").tolist()
labels = predictions.get_field("labels").tolist()
labels = [coco_demo.CATEGORIES[i] for i in labels]
boxes = predictions.bbox
rects = []
import torch
for box, score, label in zip(boxes, scores, labels):
box = box.to(torch.int64)
top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
r = [top_left[0], top_left[1], bottom_right[0], bottom_right[1]]
rect = {'class': label, 'conf': score, 'rect': r}
rects.append(rect)
import numpy as np
im_mask = np.copy(im)
draw_bb(im_mask, [r['rect'] for r in rects], [r['class'] for r in rects],
[r['conf'] for r in rects])
import json
from process_image import show_images
show_images([im, im_mask], 1, 2)
def build_predictor(self, confidence_threshold, car):
if car:
pred = COCODemo_car(
self.cfg,
min_image_size=1080,
confidence_threshold=confidence_threshold,
)
else:
pred = COCODemo(
self.cfg,
min_image_size=1080,
confidence_threshold=confidence_threshold,
)
self.pred = pred
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Inference")
parser.add_argument('--config', default='/output/tf_dir/config.yml')
parser.add_argument('--model_weights',
default='/output/tf_dir/model_final.pth')
parser.add_argument('--input_root', default='/input0/train_0712/images/')
parser.add_argument('--output_dir', default='./predictions')
args = parser.parse_args()
# update the config options with the config file
cfg.merge_from_file(args.config)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
coco_demo = COCODemo(cfg,
min_image_size=800,
confidence_threshold=0.7,
weight_loading=args.model_weights)
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
for img_name in tqdm(os.listdir(args.input_root)[:10]):
img = load(os.path.join(args.input_root, img_name))
predictions = coco_demo.run_on_opencv_image(img)
cv2.imwrite('%s/%s' % (args.output_dir, img_name), predictions)
def remove_background(input_path):
config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
image = cv2.imread(input_path, cv2.IMREAD_COLOR)
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
)
predictions, isfind = coco_demo.run_on_opencv_image(image,
category=3,
is_crop=True)
img = resize2SquareKeepingAspectRation(predictions, 256, cv2.INTER_AREA)
final_path = "temp_no_background.png"
cv2.imwrite(final_path, img)
return final_path
def initialize_model():
# Initialize the model
# args: None
# return: [boolean]: Error loading model
config_file = f"./configs/e2e_faster_rcnn_{server.architecture}.yaml"
print(f'trying to load model from {config_file}')
if not os.path.exists(config_file):
print("Dir does not exists")
return True
cfg.merge_from_file(config_file)
# os.path.join(args.DATA_FOLDER, args.WEIGHT)
cfg.DATASETS.DATA_DIR = os.path.join(args.DATA_FOLDER, args.PROJECT_FOLDER)
cfg.DATASETS.TRAIN = ['folder_train']
cfg.DATASETS.TEST = ['folder_test']
cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 3
cfg.TEST.DETECTIONS_PER_IMG = 200
cfg.MODEL.ROI_HEADS.DETECTIONS_PER_IMG = 200
# cfg.RPN.NMS_THRESH = 0.7
cfg.MODEL.ROI_HEADS.NMS = 0.1 #0.5
# cfg.TEST.BBOX_AUG.ENABLED=True
# cfg.TEST.BBOX_AUG.H_FLIP=True
# cfg.OUTPUT_DIR = os.path.dirname(cfg.DATASETS.DATA_DIR)
cfg.OUTPUT_DIR = os.path.join(
# cfg.DATASETS.DATA_DIR, "output", server.architecture)
cfg.DATASETS.DATA_DIR,
"output",
args.WEIGHT)
# cfg.freeze()
# print(cfg.OUTPUT_DIR + "/model_bestval.pth")
server.model = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
weight_loading=cfg.OUTPUT_DIR + "/model_bestval.pth",
)
# Get list of CATEGORIES from dataloader
data_loader = make_data_loader(cfg, is_train=False)
server.model.CATEGORIES = data_loader[0].dataset.CLASSES
return False
def __init__(
self,
region_type,
ground_level,
ground_removal_th,
use_maskrcnn,
city_name,
min_point_num,
dbscan_eps,
calib_data,
path_debug_output,
ground_removal_method,
dataset_name,
path_rcnn_config="../mask_rcnn/models/e2e_mask_rcnn_R_50_FPN_1x.yaml",
):
self.region_type = region_type
self.ground_level = ground_level
self.ground_removal_th = ground_removal_th
self.use_maskrcnn = use_maskrcnn
self.city_name = city_name
self.min_point_num = min_point_num
self.dbscan_eps = dbscan_eps
self.calib_data = calib_data
self.path_debug_output = path_debug_output
self.ground_removal_method = ground_removal_method
self.dataset_name = dataset_name
if use_maskrcnn:
cfg.merge_from_file(path_rcnn_config)
cfg.freeze()
# using setting in official demo
coco_demo_detector = COCODemo(
cfg,
confidence_threshold=0.7,
min_image_size=224,
)
self.mask_rcnn_detector = coco_demo_detector
def main():
#get the root of the packacge
rospack = rospkg.RosPack()
package_root = rospack.get_path('maskrcnnpkg')
rospy.init_node('image_editor')
rospy.loginfo("in main")
# Define your image topic
#create the setup the maskrcnn
cfg.merge_from_file(package_root+"/scripts/pre_trained_models/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml")
cfg.merge_from_list(['MODEL.DEVICE', 'cpu'])
cfg.freeze()
coco_demo = COCODemo(
cfg,
confidence_threshold=0.7,
masks_per_dim=2,
min_image_size=250,
)
# cam = cv2.VideoCapture(1)
# ret_val, img = cam.read()
# composite = coco_demo.run_on_opencv_image(img)
# cv2.imshow("processing", composite)
# cv2.waitKey(5000)
# cv2.destroyAllWindows()
image_topic = "/camera_end_effector/image_raw"
edited_topic = "/image_processed"
# Set up your subscriber and define its callback
pub = rospy.Publisher(edited_topic, Image, queue_size=10)
sub = rospy.Subscriber(image_topic, Image, display_image,callback_args = [pub,coco_demo])
#pub.publish(msg_frame)
#rospy.Subscriber(stream_image_topic, Image, image_callback, callback_args=pub)
# Spin until ctrl + c
rospy.loginfo("subscribed")
rospy.spin()
def main():
parser = argparse.ArgumentParser(description="PyTorch Object Detection Webcam Demo")
parser.add_argument(
"--config-file",
default="../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument(
"--confidence-threshold",
type=float,
default=0.7,
help="Minimum score for the prediction to be shown",
)
parser.add_argument(
"--min-image-size",
type=int,
default=224,
help="Smallest size of the image to feed to the model. "
"Model was trained with 800, which gives best results",
)
parser.add_argument(
"--show-mask-heatmaps",
dest="show_mask_heatmaps",
help="Show a heatmap probability for the top masks-per-dim masks",
action="store_true",
)
parser.add_argument(
"--masks-per-dim",
type=int,
default=2,
help="Number of heatmaps per dimension to show",
)
parser.add_argument(
"opts",
help="Modify model config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
# load config from file and command-line arguments
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
# prepare object that handles inference plus adds predictions on top of image
coco_demo = COCODemo(
cfg,
confidence_threshold=args.confidence_threshold,
show_mask_heatmaps=args.show_mask_heatmaps,
masks_per_dim=args.masks_per_dim,
min_image_size=args.min_image_size,
)
cam = cv2.VideoCapture(0)
while True:
start_time = time.time()
ret_val, img = cam.read()
composite = coco_demo.run_on_opencv_image(img)
print("Time: {:.2f} s / img".format(time.time() - start_time))
cv2.imshow("COCO detections", composite)
if cv2.waitKey(1) == 27:
break # esc to quit
cv2.destroyAllWindows()
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Webcam Demo")
parser.add_argument(
"--config-file",
default="configs/test.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument(
"--confidence-threshold",
type=float,
default=0.7,
help="Minimum score for the prediction to be shown",
)
parser.add_argument(
"--min-image-size",
type=int,
default=224,
help="Smallest size of the image to feed to the model. "
"Model was trained with 800, which gives best results",
)
parser.add_argument(
"--show-mask-heatmaps",
dest="show_mask_heatmaps",
help="Show a heatmap probability for the top masks-per-dim masks",
action="store_true",
)
parser.add_argument(
"--masks-per-dim",
type=int,
default=2,
help="Number of heatmaps per dimension to show",
)
parser.add_argument(
"opts",
help="Modify model config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
# load config from file and command-line arguments
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
# prepare object that handles inference plus adds predictions on top of image
coco_demo = COCODemo(
cfg,
confidence_threshold=args.confidence_threshold,
show_mask_heatmaps=args.show_mask_heatmaps,
masks_per_dim=args.masks_per_dim,
min_image_size=args.min_image_size,
)
cam = cv2.VideoCapture(0)
while True:
start_time = time.time()
ret_val, img = cam.read()
composite = coco_demo.run_on_opencv_image(img)
print("Time: {:.2f} s / img".format(time.time() - start_time))
cv2.imshow("COCO detections", composite)
if cv2.waitKey(1) == 27:
break # esc to quit
cv2.destroyAllWindows()
image = np.array(pil_image)[:, :, [2, 1, 0]]
return image
def imshow(img):
plt.imshow(img[:, :, [2, 1, 0]])
plt.axis("off")
config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
config_file = "../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cpu"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
)
# load image and then run prediction
#image = load("http://farm3.staticflickr.com/2469/3915380994_2e611b1779_z.jpg")
image = Image.open('/media/HDD_4TB/MSCOCO/images/train2017/000000579043.jpg')
image = np.array(image)[:, :, [2, 1, 0]]
predictions = coco_demo.run_on_opencv_image(image)
imshow(predictions)
class MaskRcnnRos(): #RosCppCommunicator):
def __init__(
self,
config_path=os.path.dirname(__file__) +
"/configs/caffe2/e2e_mask_rcnn_X_101_32x8d_FPN_1x_caffe2.yaml"):
self.model_config_path = config_path
cfg.merge_from_file(self.model_config_path)
cfg.merge_from_list([])
cfg.freeze()
self._greyscale_palette = (2 * 25 - 1)
# prepare object that handles inference plus adds predictions on top of image
self.coco_demo = COCODemo(cfg,
confidence_threshold=0.8,
show_mask_heatmaps=False,
masks_per_dim=1,
min_image_size=800)
self._greyscale_colours = self._generate_grayscale_values()
self.mask_rcnn_service = rospy.Service("MaskRcnnService", MaskRcnn,
self.mask_rcnn_service_callback)
print("Ready to receive mask rcnn service calls.")
self.mask_rcnn_publisher = rospy.Publisher('mask_rcnn_img',
Image,
queue_size=10)
@torch.no_grad()
def mask_rcnn_service_callback(self, req):
response = MaskRcnnResponse()
input_image = ros_numpy.numpify(req.input_image)
response_image, labels, label_indexs = self.analyse_image(input_image)
display_image = response_image * 48
output_image_msg = ros_numpy.msgify(Image,
response_image,
encoding='mono8')
display_image_msg = ros_numpy.msgify(Image,
display_image,
encoding='mono8')
self.mask_rcnn_publisher.publish(display_image_msg)
response.success = True
response.output_mask = output_image_msg
response.labels = labels
response.label_indexs = label_indexs
del response_image
del labels
del label_indexs
return response
@torch.no_grad()
def analyse_image(self, image):
predictions = self.coco_demo.compute_prediction(image)
top_predictions = self.coco_demo.select_top_predictions(predictions)
return self.create_pixel_masks(image, top_predictions)
def create_pixel_masks(self, image, predictions):
"""
Adds the instances contours for each predicted object.
Each label has a different color.
Arguments:
image (np.ndarray): an image as returned by OpenCV
predictions (BoxList): the result of the computation by the model.
It should contain the field `mask` and `labels`.
"""
width = image.shape[0]
height = image.shape[1]
blank_mask = np.zeros((width, height), np.uint8)
if predictions is None:
return blank_mask, [], []
masks = predictions.get_field("mask")
label_indexs = predictions.get_field("labels").numpy()
labels = self.convert_label_index_to_string(label_indexs)
# colours = self.get_greyscale_colours(label_indexs)
if masks.ndim < 3:
masks = np.expand_dims(masks, axis=0)
masks = np.expand_dims(masks, axis=0)
#TODO: make sure there is a boarder around each mask so that they are definetely considered
#separate objects
for mask, semantic_index in zip(masks, label_indexs):
thresh = mask[0, :, :].astype(np.uint8) * semantic_index
# print(mask.shape)
# thresh = mask.astype(np.uint8) * colour
# print(thresh)
blank_mask += thresh
# contours, hierarchy = cv2_util.findContours(
# thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
# )
# # contours = contours[0] if len(contours) == 2 else contours[1]
# image = cv2.drawContours(image, contours, -1, color, 3)
composite = blank_mask
return composite, labels, label_indexs
def convert_label_index_to_string(self, labels):
return [self.coco_demo.CATEGORIES[i] for i in labels]
def get_single_label_from_index(self, label):
return self.coco_demo.CATEGORIES[label]
def get_greyscale_colours(self, label_index):
return self._greyscale_colours[label_index]
def _generate_grayscale_values(self):
"""[Generates n number of distinct values between 1 and 255 for each label. This should be
used for visualisation purposes only as VDOSLAM just needs a distinct value]
Returns:
[List]: [List of values]
"""
numer_of_cats = len(self.coco_demo.CATEGORIES)
categories_index = np.linspace(0, numer_of_cats, numer_of_cats + 1)
colors = np.array(categories_index) * self._greyscale_palette
colors = (colors % 255).astype("uint8")
print(type(colors))
return colors
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Webcam Demo")
parser.add_argument(
"--config-file",
default="../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument(
"--confidence-threshold",
type=float,
default=0.7,
help="Minimum score for the prediction to be shown",
)
parser.add_argument(
"--min-image-size",
type=int,
default=224,
help="Smallest size of the image to feed to the model. "
"Model was trained with 800, which gives best results",
)
parser.add_argument(
"--show-mask-heatmaps",
dest="show_mask_heatmaps",
help="Show a heatmap probability for the top masks-per-dim masks",
action="store_true",
)
parser.add_argument(
"--masks-per-dim",
type=int,
default=2,
help="Number of heatmaps per dimension to show",
)
parser.add_argument(
"opts",
help="Modify model config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
# load config from file and command-line arguments
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
# prepare object that handles inference plus adds predictions on top of image
# print(args.min_image_size)
coco_demo = COCODemo(
cfg,
confidence_threshold=args.confidence_threshold,
show_mask_heatmaps=args.show_mask_heatmaps,
masks_per_dim=args.masks_per_dim,
# min_image_size=args.min_image_size,
)
mask_thread = Thread(target=mask_in_thread, args=(coco_demo, ))
mask_thread.start()
cam = cv2.VideoCapture(0)
cv2_win = 'MASKRCNN (COCO)'
cv2.namedWindow(cv2_win, cv2.WINDOW_NORMAL)
# time_cma = None
# n = 0
mask_count = 0
while True:
ret_val, img = cam.read()
mutex.acquire()
mask_thread_dict['frame'] = img
mutex.release()
# start_time = time.time()
# composite = coco_demo.run_on_opencv_image(img)
# elapsed = time.time() - start_time
# if time_cma:
# time_cma = time_cma + (elapsed - time_cma) / (n+1)
# else:
# time_cma = elapsed
composite = img.copy()
mutex.acquire()
top_predictions = mask_thread_dict['top_predictions']
mutex.release()
if top_predictions:
if coco_demo.cfg.MODEL.MASK_ON:
masks = coco_demo.crop_mask_only(composite, top_predictions)
for mask in masks:
cv2.imwrite('masks/{}.png'.format(mask_count), mask)
mask_count += 1
# composite = coco_demo.overlay_mask( composite, top_predictions )
# composite = coco_demo.overlay_boxes( composite, top_predictions )
elif coco_demo.show_mask_heatmaps:
composite = coco_demo.create_mask_montage(
composite, top_predictions)
# print("Time: {:.2f} s / img".format(time.time() - start_time))
cv2.imshow(cv2_win, composite)
if cv2.waitKey(1) == 27:
break # esc to quit
# n += 1
cv2.destroyAllWindows()
mask_thread_dict['keep_going'] = False
def main():
parser = argparse.ArgumentParser(
description="PyTorch Object Detection Webcam Demo")
parser.add_argument(
"--config-file",
default=
"./configs/yolonet/yolonet_mask_R-101-FPN_2x_adjust_std011_ms.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument(
"--confidence-threshold",
type=float,
default=0.5,
help="Minimum score for the prediction to be shown",
)
parser.add_argument(
"--min-image-size",
type=int,
default=416,
help="Smallest size of the image to feed to the model. "
"Model was trained with 800, which gives best results",
)
parser.add_argument(
"--show-mask-heatmaps",
dest="show_mask_heatmaps",
help="Show a heatmap probability for the top masks-per-dim masks",
action="store_true",
)
parser.add_argument(
"--masks-per-dim",
type=int,
default=2,
help="Number of heatmaps per dimension to show",
)
parser.add_argument(
"opts",
help="Modify model config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
# load config from file and command-line arguments
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
# prepare object that handles inference plus adds predictions on top of image
coco_demo = COCODemo(
cfg,
confidence_threshold=args.confidence_threshold,
show_mask_heatmaps=args.show_mask_heatmaps,
masks_per_dim=args.masks_per_dim,
min_image_size=args.min_image_size,
)
files = os.listdir(
"/home/stirfryrabbit/Projects/Research_Project/sheepCount/retinamask/train_test_data/coco-style/val2017"
)
# print(files)
# cam = cv2.VideoCapture(0)
# while True:
for file in files:
img = cv2.imread(
"/home/stirfryrabbit/Projects/Research_Project/sheepCount/retinamask/train_test_data/coco-style/val2017/{}"
.format(file))
start_time = time.time()
# ret_val, img = cam.read()
composite = coco_demo.run_on_opencv_image(img)
print("Time: {:.2f} s / img".format(time.time() - start_time))
cv2.imshow("COCO detections", composite)
if cv2.waitKey(1000) == 27:
break # esc to quit
cv2.destroyAllWindows()
project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# cfg.merge_from_file(
# os.path.join(project_dir,
# "configs/caffe2/e2e_faster_rcnn_R_50_FPN_1x_caffe2.yaml"))
cfg.merge_from_file('../configs/e2e_faster_rcnn_R_50_FPN_1x.yaml')
cfg.MODEL.WEIGHT = '../weights/model_final.pth'
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
cfg.freeze()
logging.info('converting model weights from: {}'.format(cfg.MODEL.WEIGHT))
# prepare object that handles inference plus adds predictions on top of image
coco_demo = COCODemo(
cfg,
confidence_threshold=0.7,
show_mask_heatmaps=False,
masks_per_dim=2,
min_image_size=480,
)
# prepare for onnx export
coco_demo.model.prepare_onnx_export()
def single_image_to_top_predictions(image):
image_list = ImageList(image.unsqueeze(0), [(int(image.size(-2)), int(image.size(-1)))])
for param in coco_demo.model.parameters():
param.requires_grad = False
result, = coco_demo.model(image_list)
scores = result.get_field("scores")
sub_dict['result'].append(image_dict)
with open("record.json", "w") as f:
json.dump(sub_dict, f)
print("加载入文件完成...")
if __name__ == "__main__":
config_file = "../configs/e2e_faster_rcnn_X_101_32x8d_FPN_1x_val.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
)
# print(model.summary())
# load label to names mapping for visualization purposes
labels_to_names = {
0: 'tieke',
1: 'heiding',
2: 'daoju',
3: 'dian',
4: 'jiandao'
}
test_img_fold = '/home/hywel/Documents/keras-retinanet/keras_retinanet/CSV/jinnan2_round1_train_20190305/restricted'
for file_name in model_list:
# config_file = "../configs/caffe2/e2e_mask_rcnn_X-152-32x8d-FPN-IN5k_1.44x_caffe2.yaml"
config_file = "../configs/CVPR_model_tests/" + file_name
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
confidence_threshold = 0.5
min_image_size = 600
coco_demo = COCODemo(
cfg,
min_image_size=min_image_size,
confidence_threshold=confidence_threshold,
)
def load(image_path):
"""
Given an url of an image, downloads the image and
returns a PIL image
"""
pil_image = Image.open(image_path).convert("RGB")
# convert to BGR format
image = np.array(pil_image)[:, :, [2, 1, 0]]
return image
def do_detection(image_path):
image = load(image_path)
