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)