使用预测框坐标裁剪图像

Crop image using pred boxes coordinates

我使用 detectron2 来预测物体在图像中的位置。现在我正在尝试使用预测框来裁剪图像(在我的用例中,每张图像仅检测到 1 object/box)。与我的问题相关的代码部分如下。问题是它只裁剪图像的左侧,但我需要它(显然)也裁剪顶部、右侧和底部,以便裁剪到检测到的对象的形状。原始图像的形状为 (x, y, 3),因此它们是 RGB 图像。我错过了什么?

from detectron2.utils.visualizer import ColorMode
import glob

imageName = "my_img.jpg"
im = cv2.imread(imageName)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1], metadata=test_metadata, scale=0.8)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])

boxes = outputs["instances"].pred_boxes
boxes = list(boxes)[0].detach().cpu().numpy()

# extract the bounding box coordinates
(x, y) = (int(boxes[0]), int(boxes[1]))
(w, h) = (int(boxes[2]), int(boxes[3]))
crop_img = image[x:y+h, y:x+w]
cv2_imshow(crop_img)

我也尝试了以下方法,但它 trim 从顶部拍摄了太多图像,根本没有 trim 图像的右侧或底部。

from detectron2.data.transforms import CropTransform

ct = CropTransform(x, y, w, h)
crop_img = ct.apply_image(image)
cv2_imshow(crop_img)

尝试使用它,我能够使用以下命令在检测到的框周围裁剪图像,但它并不理想,因为我必须对其进行硬编码。

crop_img = image[y-40:y+h-390, x:x+w-395]

以下应该有效。

def crop_object(image, box):
  """Crops an object in an image

  Inputs:
    image: PIL image
    box: one box from Detectron2 pred_boxes
  """

  x_top_left = box[0]
  y_top_left = box[1]
  x_bottom_right = box[2]
  y_bottom_right = box[3]
  x_center = (x_top_left + x_bottom_right) / 2
  y_center = (y_top_left + y_bottom_right) / 2

  crop_img = image.crop((int(x_top_left), int(y_top_left), int(x_bottom_right), int(y_bottom_right)))
  return crop_img

# Get pred_boxes from Detectron2 prediction outputs
boxes = outputs["instances"].pred_boxes
# Select 1 box:
box = list(boxes)[0].detach().cpu().numpy()
# Crop the PIL image using predicted box coordinates
crop_img = crop_object(image, box)