WIP: process container masks and compute principal axis ratios

MaskRCNN_classify.py

@@ -0,0 +1,356 @@
+# This is a script that loads the generic container detection model and applies it to an image. This script
+# was developed purely to apply the model to the 55,000 image dataset which included container labels but
+# did not include container segmentations. The model is therefore applied to these images in order to create
+# a json file for this dataset that contains both the segmentation masks as well as the container label in
+# coco format - the format required for training Mask RCNN on onepanel.
+
+import sys
+import numpy as np
+import pandas as pd
+import cv2
+import glob
+from datetime import datetime
+import logging
+import time
+
+# Import mrcnn libraries
+from mrcnn.config import Config
+from mrcnn import visualize
+import mrcnn.model as modellib
+import os
+import skimage
+import random
+
+from samples.coco import coco
+import csv
+
+class_names = ['other', 'container']
+
+# downloaded the new model here
+WEIGHTS_DIR = '/home/kunal/new_ssd/FF_AWS_Stuff/models_n_classes/onepanel_trained_model.h5'
+
+# default folder for output images
+IMG_FOLDER = "out/images"
+
+ROOT_DIR = os.path.join(os.getcwd(), "mrcnn")
+assert os.path.exists(
+    ROOT_DIR), 'ROOT_DIR does not exist. Did you forget to read the instructions above? ;)'
+sys.path.append(ROOT_DIR)
+
+# Directory to save logs and trained model
+MODEL_DIR = os.path.join(ROOT_DIR, "logs")
+
+
+class InferenceConfig(coco.CocoConfig):
+    # Set batch size to 1 since we'll be running inference on
+    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
+    NUM_CLASSES = 2
+    GPU_COUNT = 1
+    IMAGES_PER_GPU = 1
+
+
+def display_image(image, resize):
+    """Display an Image with optional resizing
+     User needs to press ESC key to close image window
+     and proceed further
+    """
+    if resize:
+        dim = image.shape
+        height = dim[0]
+        width = dim[1]
+        # resize the output image before display
+        smaller_image = cv2.resize(image, (int(width / 2), int(height / 2)))
+        # display the image & wait for esc key
+        cv2.imshow('Image', smaller_image)
+    else:
+        cv2.imshow('Image', image)
+    cv2.waitKey()
+    cv2.destroyAllWindows()
+
+
+def get_gray_img_for_mask(mask):
+    """Returns a binary image from given Mask"""
+
+    # get image dimensions
+    height, width, channels = mask.shape
+    if channels == 0:
+        # generate a zero-array
+        integer_mask = numpy.zeros(width * height, dtype=np.uint8)
+    else:
+        # flatten the mask to a 1-d array
+        flat_mask = mask.flatten()
+        integer_mask = flat_mask.astype(np.uint8)
+    # re-size to image dimensions, we get 1-ch image
+    gray_mask = integer_mask.reshape((height, width, 1))
+    return gray_mask
+
+
+def get_color_img_for_mask(mask):
+    """Retuns a color image from given mask"""
+
+    gray_mask = get_gray_img_for_mask(mask)
+    mask_rgb = cv2.cvtColor(gray_mask, cv2.COLOR_GRAY2RGB)
+    return mask_rgb
+
+
+def get_contours_from_mask(mask):
+    """returns CV Contours from a given mask"""
+    gray_img = get_gray_img_for_mask(mask)
+    gray_img = gray_img * 255
+    # display_image(gray_img, True)
+    contours, hierarchy = cv2.findContours(
+        gray_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    return contours
+
+
+def get_masked_image(image, mask):
+    """Returns a masked version of the given image"""
+
+    # 1. get a gray_img from given mask
+    gray_mask = get_gray_img_for_mask(mask)
+    # convert gray mask to rgb
+    mask_rgb = cv2.cvtColor(gray_mask, cv2.COLOR_GRAY2RGB)
+    # multiply mask with test image to get output image
+    image = image * mask_rgb
+    enable_display = False
+    if enable_display:
+        display_image(image, True)
+    return image
+
+
+def compute_rectangles(contoured_image, contours):
+    """Returns array of principle rectangles in the
+    given image based on given contours,
+    Also draws the rectangles on given image"""
+
+    # for each contour compute principle
+    # axis ratios, finally we use the ratio
+    # from largest object
+    rect_area = 0
+    principle_rect = [[0, 0], [0, 0], 0]
+    index = 0
+    for contour in contours:
+        rect = cv2.minAreaRect(contour)
+        # rect is (center(x, y), (width, height), angle of rotation)
+        w, h = rect[1]
+        area = h * w
+        # largest object
+        if area > rect_area:
+            rect_area = area
+            principle_rect = rect
+        box = cv2.boxPoints(rect)
+        box = np.int0(box)
+        cv2.drawContours(contoured_image, [box], -1, (0, 0, 255), 4)
+        index = index + 1
+
+    # display_image(contoured_image,True)
+    return principle_rect
+
+
+def get_principle_rect_from_image(contoured_image, mask):
+    """Return array of principle rectangles for
+    object in an image, based on the given mask
+    """
+
+    # 1. get a gray_img from given mask
+    gray_mask = get_gray_img_for_mask(mask)
+
+    # 2. get contours from given mask
+    contours = get_contours_from_mask(mask)
+
+    # 3. draw contours
+    cv2.drawContours(contoured_image, contours, -1, (0, 255, 0), 3)
+
+    # 4. draw rectangles & get principle rect data
+    principle_rect = compute_rectangles(contoured_image, contours)
+
+    return principle_rect
+
+
+def process_classification_results(filename, mask, logger):
+    """Further process results from mask-RCNN"""
+
+    # read the test image
+    test_image = cv2.imread(filename, 1)
+    # create a copy for drawing contours
+    contoured_image = test_image.copy()
+    # create a copy for creating masked image
+    masked_image = test_image.copy()
+
+    # get image dimensions
+    dimensions = test_image.shape
+    height = dimensions[0]
+    width = dimensions[1]
+    channels = dimensions[2]
+
+    # iterate through each mask
+    for ch in range(0, mask.shape[2]):
+
+        # process each mask
+        logger.debug("processing channel num: %d", ch)
+        sub_mask = mask[:, :, ch]
+        sub_mask = sub_mask.reshape((sub_mask.shape[0], sub_mask.shape[1], 1))
+
+        # 1. process image & mask to get principle rect
+        principle_rect = get_principle_rect_from_image(
+            contoured_image, sub_mask)
+
+        # 2. accumulate multiple masks, for future reference
+        if ch == 0:
+            cum_mask = sub_mask
+        else:
+            cum_mask = cum_mask + sub_mask
+
+        # 3. all done, log the rectangle data
+        logger.debug("Ch# %d, rectangle data is %s", ch, principle_rect)
+
+    masked_image = get_masked_image(masked_image, cum_mask)
+    # display_image(masked_image,True)
+
+    # create output file-path with a post-fix
+    head, tail = os.path.split(filename)
+    splitname = os.path.splitext(tail)
+    out_filename = splitname[0] + "_processed" + ".jpg"
+    out_filepath = os.path.join(IMG_FOLDER, out_filename)
+
+    # concatenate masked image with contoured image
+    concatenate_output = True
+    if concatenate_output:
+        out_image = cv2.hconcat([contoured_image, masked_image])
+
+    # write the output image
+    cv2.imwrite(out_filepath, out_image)
+
+
+def createLogger():
+    """creates a logger object"""
+
+    # create log folder if it does not exist
+    # if not os.path.exists("out"):
+    #    os.mkdir("out")
+
+    log_dir = os.path.join(os.getcwd(), "out")
+    old_log_files = glob.glob(log_dir + '/*.log')
+
+    now = datetime.now()
+    dt_string = now.strftime("%d_%m_%Y_%H-%M-%S")
+    log_file_name = "log_" + dt_string + ".log"
+    log_file_name = os.path.join(log_dir, log_file_name)
+
+    logging.basicConfig(filename=log_file_name)
+    # logging.basicConfig()
+
+    logger = logging.getLogger('3DQ_Tester')
+    logger.setLevel(logging.DEBUG)
+
+    return logger
+
+    # remove old logs
+    # for f in old_log_files:
+    #    if os.path.exists(f):
+    #        try:
+    #            os.remove(f)
+    #        except OSError as e:
+    # self.__aivero_logger.error("Error: can not remove %s : %s", f,
+    # e.strerror)
+
+
+def classify_single_image(filename, model, logger):
+    """classivy an image with mask-RCNN
+    and post-process the classifier results
+    """
+
+    logger.debug("... filename is %s", filename)
+    image = skimage.io.imread(filename)
+    result = model.detect([image], verbose=1)
+    logger.debug("... detection has been done")
+
+    rs = result[0]
+
+    # this is giving error
+    # visualize.display_instances(
+    #        image, rs['rois'], rs['masks'], rs['class_ids'],
+    #        class_names, rs['scores'],
+    #        show_bbox=False, show_mask=False,
+    #        title="Predictions")
+
+    mask = rs['masks']
+    # post-process non-empty masks
+    # we can have multiple masks
+    if mask.shape[2] > 0:
+        process_classification_results(filename, mask, logger)
+    else:
+        logger.debug("there is no Mask data for %s  ********** ", filename)
+
+
+def classify_multiple_images(test_data_path, model, logger):
+    """classify all images from a given folder
+    with mask-RCNN and post-process
+    the classifier results
+    """
+
+    # process all images under test_data_path
+    for count, jpegfile in enumerate(os.listdir(test_data_path)):
+        splitname = os.path.splitext(jpegfile)
+        if splitname[1] == ".jpg":
+            # we only want to process the jpg files
+            logger.debug("processing image number %d", count)
+            filename = os.path.join(test_data_path, jpegfile)
+            t1_msec = int(round(time.time() * 1000))
+            classify_single_image(filename, model, logger)
+            t2_msec = int(round(time.time() * 1000))
+            total_time = t2_msec - t1_msec
+            logger.debug("processing done in %d seconds", total_time)
+            logger.debug("******")
+
+
+def main(args):
+    """Main function for classifying images
+    from a given folder
+    """
+
+    # crate logger object
+    logger = createLogger()
+
+    if len(args) != 2:
+        print(f"wrong usage, try: \"python MaskRCNN_classify.py <image-file-or-folder>\"")
+        logger.debug("wrong usage, try: \"python MaskRCNN_classify.py <image-file-or-folder>\"")
+        sys.exit()
+
+    test_data_path = args[1]
+    logger.debug("Process data from : %s", test_data_path)
+
+    # create out folder if it does not exist
+    if not os.path.exists("out"):
+        os.mkdir("out")
+
+    # update IMG_FOLDER name
+    now = datetime.now()
+    dt_string = now.strftime("%d_%m_%Y_%H-%M-%S")
+    global IMG_FOLDER
+    IMG_FOLDER = os.path.join("out", dt_string)
+    logger.debug("the out folder is: %s", IMG_FOLDER)
+    # create image folder if it does not exist
+    if not os.path.exists(IMG_FOLDER):
+        os.mkdir(IMG_FOLDER)
+
+    # setup classifier
+    config = InferenceConfig()
+    config.display()
+    # Create model object in inference mode.
+    model = modellib.MaskRCNN(
+        mode="inference",
+        config=config,
+        model_dir=MODEL_DIR)
+    # Load weights trained on MS-COCO
+    model.load_weights(WEIGHTS_DIR, by_name=True)
+    # , exclude=["mrcnn_bbox_fc", "mrcnn_class_logits", "mrcnn_mask "])
+
+    if os.path.isdir(test_data_path):
+        classify_multiple_images(test_data_path, model, logger)
+    else:
+        classify_single_image(test_data_path, model, logger)
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv))

Pipfile

@@ -0,0 +1,24 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+numpy = "*"
+scipy = "*"
+matplotlib = "*"
+scikit-image = "*"
+tensorflow = ">=2.0.0"
+opencv-python = "*"
+h5py = "*"
+imgaug = "*"
+ipython = {extras = ["all"], version = "*"}
+Pillow = "*"
+Cython = "*"
+pandas = "*"
+pycocotools = "*"
+
+[dev-packages]
+
+[requires]
+python_version = "3.8"