detectron2 crop results

crop_cell.py

import numpy as np
import matplotlib.pyplot as plt

def get_bboxes_from(output):
  """ returns list of bboxes """
  return outputs["instances"].__dict__['_fields']['pred_boxes'].__dict__['tensor']


def crop(bbox, in_img: np.ndarray):
  """ bbox is a list with xmin, ymin, xmax, ymax """
  xmin, ymin, xmax, ymax = bbox
  cropped_im = in_img[xmin:xmax, ymin:ymax]
  return cropped_im


def disp(im):
  plt.imshow(im)
  plt.show()
  plt.close()



def run_easy_ocr(output, im):
  
  bboxes = get_bboxes_from(output)
  for bbox in bboxes:

    crop_im = crop(bbox, in_img=im)
    # display cropped image
    disp(crop_im)

    # detect easy ocr on `crop_im` here
    # ....
    # ...
    # ..

inference_cell.py

from detectron2.utils.visualizer import ColorMode
from google.colab.patches import cv2_imshow

for thresh in [0.3, 0.5, 0.65]:

  # load predictor
  cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = thresh   # set a custom testing threshold
  predictor = DefaultPredictor(cfg)

  for dir_type in ['train', 'test']:

      # save inference models
      dataset_dicts = get_building_dicts(f"dataset/{dir_type}")
      for d in random.sample(dataset_dicts, len(dataset_dicts)):    
          im = cv2.imread(d["file_name"])
          outputs = predictor(im)  # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
          run_easy_ocr(output, im)
          
          v = Visualizer(im[:, :, ::-1],
                        metadata=building_metadata, 
                        scale=0.5, 
                        instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models
          )
          out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
          im = out.get_image()[:, :, ::-1]
          
          #cv2_imshow(im)
          out_dir = f'{dir_type}'
          os.makedirs(f'output_{out_dir}_T{thresh}', exist_ok=True)
          print('writing', f'output_{out_dir}_T{thresh}/{d["file_name"].split("/")[-1]}')
          cv2.imwrite(f'./output_{out_dir}_T{thresh}/{d["file_name"].split("/")[-1]}', im)