generate roidb with rpn

gen_rpn_app.py

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

################################################################
#
# This script takes an image as input and outputs the
# region proposals 
#
# Author: Aniket Adnaik ([email protected])
################################################################
# Copyright (c) 2017-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
###########################################################################
"""Generates region proposal file using mask rcnn and results 
are stroed in a pickle file.
Expects following environment variables to be set;
CAFFE2_HOME to be set to valid caffe2 distrbution
DETECTRON_HOME to be set to valid detectron distrbution
Perform inference on a single image or all images with a certain extension
(e.g., .jpg) in a folder.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from collections import defaultdict
import argparse
import cv2  # NOQA (Must import before importing caffe2 due to bug in cv2)
import glob
import logging
import os
import sys
import time
import pickle

sys.path.append(os.environ['CAFFE2_HOME'])
from caffe2.python import workspace

sys.path.append(os.environ['DETECTRON_HOME'])
from detectron.core.config import assert_and_infer_cfg
from detectron.core.config import cfg
from detectron.core.config import merge_cfg_from_file
from detectron.utils.io import cache_url
from detectron.utils.logging import setup_logging
from detectron.utils.timer import Timer
import detectron.core.test_engine as infer_engine
import detectron.datasets.dummy_datasets as dummy_datasets
import detectron.utils.c2 as c2_utils
import detectron.utils.vis as vis_utils
import pdb

c2_utils.import_detectron_ops()

# OpenCL may be enabled by default in OpenCV3; disable it because it's not
# thread safe and causes unwanted GPU memory allocations.
cv2.ocl.setUseOpenCL(False)

def gen_rpn_proposals(image_path, model):
    """gen_rpn_proposals - needed by the scene-graph generation code. Can't ground the bboxes
       if we don't have any guesses for the bboxes in the first place.
    """

    image_ext = 'jpg' 
    if os.path.isdir(image_path):
        im_list = glob.iglob(image_path + '/*.' + image_ext)
    else:
        im_list = [image_path]

    res = []
    for i, im_name in enumerate(im_list):
        im_basename = os.path.basename(im_name)
        out_name = os.path.join(
            output_dir, '{}'.format(im_basename + '.' + output_ext)
        )

        im = cv2.imread(im_name)

        with c2_utils.NamedCudaScope(0):
            cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
                model, im, None, timers=timers
            )
        res.append(cls_boxes)

    return res