Skip to content

Commit

Permalink
Add Faster R-CNN
Browse files Browse the repository at this point in the history
  • Loading branch information
Yuliang-Zou committed Aug 25, 2020
1 parent e87f286 commit d2a6b10
Show file tree
Hide file tree
Showing 40 changed files with 4,396 additions and 1 deletion.
6 changes: 6 additions & 0 deletions .gitignore
@@ -1,2 +1,8 @@
.idea/
__pycache__
_ext
*.pyc
*.so
maskrcnn_benchmark.egg-info/
build/
dist/
33 changes: 33 additions & 0 deletions configs/object_detector/e2e_faster_rcnn_R_50_FPN_1x.yaml
@@ -0,0 +1,33 @@
MODEL:
META_ARCHITECTURE: "ObjectDetRCNN"
WEIGHT: "pretrained_model/e2e_faster_rcnn_R_50_FPN_1x.pth"
BACKBONE:
CONV_BODY: "R-50-FPN"
RESNETS:
BACKBONE_OUT_CHANNELS: 256
RPN:
USE_FPN: True
ANCHOR_STRIDE: (4, 8, 16, 32, 64)
PRE_NMS_TOP_N_TRAIN: 2000
PRE_NMS_TOP_N_TEST: 1000
POST_NMS_TOP_N_TEST: 1000
FPN_POST_NMS_TOP_N_TEST: 1000
ROI_HEADS:
USE_FPN: True
ROI_BOX_HEAD:
POOLER_RESOLUTION: 7
POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
POOLER_SAMPLING_RATIO: 2
FEATURE_EXTRACTOR: "FPN2MLPFeatureExtractor"
PREDICTOR: "FPNPredictor"
NUM_CLASSES: 81
DATASETS:
TRAIN: ("coco_2014_train", "coco_2014_valminusminival")
TEST: ("coco_2014_minival",)
DATALOADER:
SIZE_DIVISIBILITY: 32
SOLVER:
BASE_LR: 0.02
WEIGHT_DECAY: 0.0001
STEPS: (60000, 80000)
MAX_ITER: 90000
Binary file added demo/coco_demo.jpg
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
77 changes: 77 additions & 0 deletions demo/demo_obj_det.py
@@ -0,0 +1,77 @@
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import argparse
import cv2

from maskrcnn_benchmark.config import cfg
from predictor import COCODemo

import time


def main():
parser = argparse.ArgumentParser(description="PyTorch Object Detection Webcam Demo")
parser.add_argument(
"--config-file",
default="configs/object_detector/e2e_faster_rcnn_R_50_FPN_1x.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument(
"--confidence-threshold",
type=float,
default=0.1,
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,
weight_loading='pretrained_model/e2e_faster_rcnn_R_50_FPN_1x.pth'
)

img = cv2.imread('demo/coco_demo.jpg')
# Get a dict of results
result_dict = coco_demo.get_result_dict(img)
# Visualization
#composite = coco_demo.run_on_opencv_image(img)
#cv2.imwrite('test.png', composite)


if __name__ == "__main__":
main()

0 comments on commit d2a6b10

Please sign in to comment.