import numpy as np
import torchvision
import PIL
import cv2


def crop(image):
    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    height, width, _ = image.shape

    # 指定要剪切的像素数量
    pixels_to_remove = 5

    # 剪切图像以去掉指定数量的像素
    image = image[pixels_to_remove:height - pixels_to_remove, pixels_to_remove:width - pixels_to_remove]
    # result = np.ones_like(image) * 255
    # 将图像转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    threshold_value = 75
    ret, binary_image = cv2.threshold(gray, threshold_value, 255, cv2.THRESH_BINARY)
    image = cv2.cvtColor(binary_image, cv2.COLOR_GRAY2RGB)

    # 进行边缘检测
    edges = cv2.Canny(binary_image, 50, 150)

    # 寻找轮廓
    contours, hierarchy = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 设置面积阈值，根据实际情况调整
    area_threshold = 200

    # 遍历轮廓
    for contour in contours:
        # 计算轮廓的面积
        area = cv2.contourArea(contour)

        # 如果面积小于阈值，则涂成白色（255）
        if area < area_threshold:
            cv2.drawContours(image, [contour], 0, (255, 255, 255), thickness=cv2.FILLED)
            # cv2.drawContours(result, [contour], 0, (0, 255, 0), thickness=2)

    # for contour in contours:
    #     # 计算轮廓的面积
    #     area = cv2.contourArea(contour)
    #
    #     # 如果面积小于阈值，则涂成白色（255）
    #     cv2.drawContours(result, [contour], 0, (0, 255, 0), thickness=1)

    # 保存处理后的图像
    # cv2.imwrite("2.png", result)
    # image = result
    # 转换为灰度图像
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 查找黑色像素的坐标
    black_pixels = np.column_stack(np.where(gray_image == 0))
    # 找到离上边框最近的黑色像素
    top_pixel = black_pixels[np.argmin(black_pixels[:, 0])]
    # 找到离下边框最近的黑色像素
    bottom_pixel = black_pixels[np.argmax(black_pixels[:, 0])]
    # 找到离左边框最近的黑色像素
    left_pixel = black_pixels[np.argmin(black_pixels[:, 1])]
    # 找到离右边框最近的黑色像素
    right_pixel = black_pixels[np.argmax(black_pixels[:, 1])]

    # 计算长方形的宽和高
    width = right_pixel[1] - left_pixel[1]
    height = bottom_pixel[0] - top_pixel[0]
    # 计算长方形的左上角坐标
    rectangle_top_left = (top_pixel[0], left_pixel[1])
    # 创建包含长方形内容的新图像
    rectangle_image = image[rectangle_top_left[0]:rectangle_top_left[0] + height,
                      rectangle_top_left[1]:rectangle_top_left[1] + width]
    # 保存新图像

    original_image = rectangle_image

    height, width, _ = original_image.shape

    # 计算正方形图片的边长，取较大的那个值
    side_length = max(width, height) + 6  # 增加6像素的边缘间隔

    # 创建一个白底的正方形图片
    square_image = np.ones((side_length, side_length, 3), dtype=np.uint8) * 255

    # 计算将原图放入正方形图片中的位置
    x_offset = (side_length - width) // 2
    y_offset = (side_length - height) // 2

    # 将原图粘贴到正方形图片的中间
    square_image[y_offset:y_offset + height, x_offset:x_offset + width] = original_image
    return square_image

def web_input(image):
    image_tmp = crop(image)
    image_transforms = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
    input_img = PIL.Image.fromarray(image_tmp)
    input_img = input_img.resize((100, 100), PIL.Image.LANCZOS)
    wd_new, ht_new = input_img.size
    wd_new = int(16 * np.ceil(wd_new / 16.0))
    ht_new = int(16 * np.ceil(ht_new / 16.0))
    input_img = input_img.resize((wd_new, ht_new), PIL.Image.LANCZOS)
    return image_transforms(input_img).unsqueeze(0), image_tmp