update for model use

README.md

@@ -260,3 +260,12 @@ note = {Accessed: [Insert date here]}
 ## License
 
 maskrcnn-benchmark is released under the MIT license. See [LICENSE](LICENSE) for additional details.
+
+## 三色图调用
+pip install git+ssh://git@git.wanmeizhensuo.com/wushengyu/maskrcnn.git
+
+```python
+from demo.three_color import handle_from_path
+path = '/xxx/xxx/xx.jpg'
+img = handle_from_path(path)        # 类似cv2.imread的img
+```

demo/three_color.py

+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+import argparse
+import cv2
+
+from maskrcnn_benchmark.config import cfg
+from demo.predictor import COCODemo
+
+
+def img_to_three_color_map(img):
+    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=800,
+        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,
+    )
+    composite = coco_demo.run_on_opencv_image(img)
+    if composite:
+        return composite
+    return None
+
+
+def handle_from_url(image_url):
+    """ 处理网路图片 """
+    cap = cv2.VideoCapture(image_url)
+    if (cap.isOpened()):
+        ret, img = cap.read()
+        return img_to_three_color_map(img)
+    return None
+
+
+def handle_from_path(image_path):
+    """ 处理本地图片 """
+    image = cv2.imread(image_path)
+    return img_to_three_color_map(image)
\ No newline at end of file