import cv2
import imageio
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

is_train = True
height, width = 384, 288
# 横 x 竖 y
bbox = [42, 40, 290, 211]  # x1,y1,x2,y2

# 1.ori img
img = imageio.imread('football.jpg')
plt.imshow(img)
plt.gca().add_patch(
    patches.Rectangle(xy=(bbox[0], bbox[1]),  # bottom, left
                      width=bbox[2] - bbox[0], height=bbox[3] - bbox[1],
                      linewidth=1, edgecolor='r', facecolor='none'))
plt.show()

add = max(img.shape[0], img.shape[1])  # 347
mean_value = [122.7717, 115.9465, 102.9801]

# 2.border img
bimg = cv2.copyMakeBorder(img,
                          add, add, add, add,
                          borderType=cv2.BORDER_CONSTANT,  # constant pixel_mean as border
                          value=mean_value)
bbox = np.array(bbox).reshape(4, ).astype(np.float32)

# bbox contains obj
objcenter = np.array([(bbox[0] + bbox[2]) / 2.,  # bbox_w/2
                      (bbox[1] + bbox[3]) / 2.])  # bbox_h/2

# shift bbox/objcenter/keypoints to new bimg
bbox += add  # move to center of bimg along with ori_img
objcenter += add

plt.imshow(bimg)
plt.gca().add_patch(
    patches.Rectangle(xy=(bbox[0], bbox[1]),  # bottom, left
                      width=bbox[2] - bbox[0], height=bbox[3] - bbox[1],
                      linewidth=1, edgecolor='r', facecolor='none'))
plt.show()

# 3.extend and crop img
bbox_extend_factor = (0.1, 0.15)

# bbox [w,h] * (1 + extend_factor), [0.1, 0.15]
crop_width = (bbox[2] - bbox[0]) * (1 + bbox_extend_factor[0] * 2)  # 两边各扩展0.1
crop_height = (bbox[3] - bbox[1]) * (1 + bbox_extend_factor[1] * 2)

if is_train:
    crop_width = crop_width * (1 + 0.25)
    crop_height = crop_height * (1 + 0.25)

print('image_wh:', img.shape[1], img.shape[0])  # 347,212
print('input_wh:', width, height)  # 288,384
print()
print('ori_bbox_wh:', bbox[2] - bbox[0], bbox[3] - bbox[1])  # 248.0,171.0
print('crop_box_wh:', crop_width, crop_height)  # 372.0,277.9
print('crop/input:', crop_width / width, crop_height / height)  # 1.29,0.72
print()

# > < depends on ori bbox size
# crop_size 取比例较大边
if crop_height / height > crop_width / width:  # height,width is model input shape (384,288)
    crop_size = crop_height
    min_shape = height
else:
    crop_size = crop_width
    min_shape = width

print('crop size:', crop_size)  # 372.0
print('min shape:', min_shape)  # 288
print()
print('after extend')
print('objcenter:', objcenter)  # 513.0,472.5
print('crop bbox:', bbox)  # [389. 387. 637. 558.]
print('bimg_wh:', bimg.shape[1], bimg.shape[0])  # 1041,906
print()

# min_shape is very important
# min_shape = height/width of input
# crop_size 与 obj 左右上下 相比较
crop_size = min(crop_size, objcenter[0] / width * min_shape * 2. - 1.)  # if width=min_shape, objcenter[0]*2-1
crop_size = min(crop_size, (bimg.shape[1] - objcenter[0]) / width * min_shape * 2. - 1)
crop_size = min(crop_size, objcenter[1] / height * min_shape * 2. - 1.)
crop_size = min(crop_size, (bimg.shape[0] - objcenter[1]) / height * min_shape * 2. - 1)

# 以 crop_size 为基准，基于 objcenter 在 bimg 上获得 左上,右下 点
# 保证图像宽高比 = model input 宽高比，所以 x,y_ratio 是相等的
min_x = int(objcenter[0] - crop_size / 2. / min_shape * width)
max_x = int(objcenter[0] + crop_size / 2. / min_shape * width)
min_y = int(objcenter[1] - crop_size / 2. / min_shape * height)
max_y = int(objcenter[1] + crop_size / 2. / min_shape * height)

x_ratio = float(width) / (max_x - min_x)
y_ratio = float(height) / (max_y - min_y)
print('ratios:', x_ratio, y_ratio)

crop_img = cv2.resize(bimg[min_y:max_y, min_x:max_x, :], (width, height))
plt.imshow(crop_img)
plt.show()

横图

ori img

border img

crop img

竖图

ori img

border img

crop img