Expand bounding boxes in edges of images to square

I have a function designed to expand bounding boxes in images to squares.

def expand_bbox(bbox, img_shape):

    h, w = bbox[2] - bbox[0], bbox[3] - bbox[1]
    if h < w:
        d = w - h
        u = min(bbox[0], d // 2) 
        v = d - u
        bbox[0] -= u
        bbox[2] += v
    else:
        d = h - w
        u = min(bbox[1], d // 2)
        v = d - u
        bbox[1] -= u
        bbox[3] += v

    bbox[0] = max(bbox[0], 0)
    bbox[1] = max(bbox[1], 0)
    bbox[2] = min(bbox[2], img_shape[0]) 
    bbox[3] = min(bbox[3], img_shape[1])
    h, w = bbox[2] - bbox[0], bbox[3] - bbox[1]
    assert h == w
    inc = h // 2
    inc = min(inc, bbox[0])
    inc = min(inc, bbox[1])
    inc = min(inc, img_shape[0] - bbox[2])
    inc = min(inc, img_shape[1] - bbox[3])
    bbox[0] -= inc
    bbox[1] -= inc
    bbox[2] += inc
    bbox[3] += inc
    h, w = bbox[2] - bbox[0], bbox[3] - bbox[1]
    assert h == w
    return bbox

It works fine for cases where the bounding box appears inside the image, but fails to expand boxes in the edges of images (i.e. when the bounding box is [979, 157, 1080, 261] and the image shape is (1080, 1920, 3))

What can I do make sure that the bounding box captures bounding boxes that appear along the edges of the images?

>Solution :

I am pretty sure all this code could be prettier, but if you search for a working code, I would go for adding some ifs:

# if increasing the upper edge position would go outside the img_shape, do not increase it, decrease bottom instead.
if bbox[2] + v > img_shape[0]:
    bbox[0] -= u + v
else:
    bbox[0] -= u
    bbox[2] += v

And the same for width.