#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File        :   roi_fix.py    
@Contact     :   zpyovo@hotmail.com
@License     :   (C)Copyright 2018-2019, Lab501-TransferLearning-SCUT
@Description :

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2024/2/27 00:02   Pengyu Zhang      1.0         None
'''

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

'''批量固定尺寸'''
import os
import cv2
from tqdm import tqdm
import numpy as np

def merge_overlapping_rectangles(rectangles):
    if len(rectangles) <= 1:
        return rectangles

    merged_rectangles = []
    for rect in sorted(rectangles, key=lambda x: x[2]*x[3], reverse=True):
        x, y, w, h = rect
        if not any((x < x2 + w2 and x + w > x2 and y < y2 + h2 and y + h > y2) for x2, y2, w2, h2 in merged_rectangles):
            merged_rectangles.append(rect)
    return merged_rectangles

def roi_Inner_detector(input,):
    # 定义结构元素（内核）
    kernel = np.ones((3, 3), np.uint8)

    # 开运算
    opening = cv2.morphologyEx(input, cv2.MORPH_OPEN, kernel)

    # 定义结构元素（内核）
    kernel = np.ones((3, 3), np.uint8)
    # 闭运算
    closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel)

    # 使用 Canny 边缘检测
    edges = cv2.Canny(closing, 50, 255, apertureSize=3)

    # 寻找轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # 临时存储矩形
    temp_rectangles = []

    # 识别矩形轮廓
    for contour in contours:
        # 轮廓近似
        epsilon = 0.001 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)

        if len(approx) >= 4:
            x, y, w, h = cv2.boundingRect(approx)
            aspect_ratio = w / float(h)
            area_contour = cv2.contourArea(contour)
            min_area_threshold = 0
            #         print('aspect_ratio',aspect_ratio)
            #         print('area_contour',area_contour)
            if aspect_ratio >= 0.80 and aspect_ratio < 1.2 and 6500 < area_contour < 13000:
                area_bounding_rect = w * h
                area_ratio = area_contour / area_bounding_rect
                # print('aspect_ratio', area_ratio)
                # print('area_ratio', area_ratio)
                # print('area_contour', area_contour)
                if 0.80 < area_ratio < 1.05:
                    temp_rectangles.append((x, y, w, h))

    # 合并重叠的方框
    merged_rectangles = merge_overlapping_rectangles(temp_rectangles)
    return merged_rectangles

input_dir = '/project/dataset/QR2024/org100/orange-0-roi-fix'
output_dir = '/project/dataset/QR2024/org100/orange-0-roi-fix2fix'

size = 128

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 遍历输入目录中的图片
for filename in tqdm(os.listdir(input_dir), desc='Processing'):
    if filename.endswith('.jpg') or filename.endswith('.png'):
        input_path = os.path.join(input_dir, filename)
        image = cv2.imread(input_path)
        height, width, _ = image.shape

        if height < size:
            image = cv2.resize(np.array(image), (size, size), interpolation=cv2.INTER_AREA)
        else:
            image = cv2.resize(np.array(image), (size, size), interpolation=cv2.INTER_LINEAR)

        image_gay = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary_image_A = cv2.threshold(image_gay, 50, 255, cv2.THRESH_BINARY)
        merged_rectangles = roi_Inner_detector(binary_image_A)

        if len(merged_rectangles) == 0 :
            _, binary_image_A = cv2.threshold(image_gay, 150, 255, cv2.THRESH_BINARY)
            merged_rectangles = roi_Inner_detector(binary_image_A)
            # print('150')
        if len(merged_rectangles) == 0 :
            _, binary_image_A = cv2.threshold(image_gay, 100, 255, cv2.THRESH_BINARY)
            merged_rectangles = roi_Inner_detector(binary_image_A)
            # print('100')
        if len(merged_rectangles) == 0 :
            _, binary_image_A = cv2.threshold(image_gay, 125, 255, cv2.THRESH_BINARY)
            merged_rectangles = roi_Inner_detector(binary_image_A)
            # print('120')
        if len(merged_rectangles) == 0 :
            _, binary_image_A = cv2.threshold(image_gay, 75, 255, cv2.THRESH_BINARY)
            merged_rectangles = roi_Inner_detector(binary_image_A)

        if len(merged_rectangles) == 0:
            _, binary_image_A = cv2.threshold(image_gay, 85, 255, cv2.THRESH_BINARY)
            merged_rectangles = roi_Inner_detector(binary_image_A)
        if len(merged_rectangles) == 0:
            _, binary_image_A = cv2.threshold(image_gay, 115, 255, cv2.THRESH_BINARY)
            merged_rectangles = roi_Inner_detector(binary_image_A)
            # print('75')
        # if len(merged_rectangles) == 0 :
        #     _, binary_image_A = cv2.threshold(image_gay, 0, 255, cv2.THRESH_BINARY)
        #     merged_rectangles = roi_Inner_detector(binary_image_A)

        if merged_rectangles:
            outermost_rectangle = max(merged_rectangles, key=lambda x: x[2] * x[3])
            x, y, w, h = outermost_rectangle
            image_crope = image[y + 1:y - 1 + h, x + 1:x - 1 + w]
            image_crope = cv2.resize(image_crope, (32, 32), interpolation=cv2.INTER_AREA)
            output_path = os.path.join(output_dir, filename)
            cv2.imwrite(output_path, image_crope)
        else:
            print("failed",filename)
        # output_path = os.path.join(output_dir, filename)
        # cv2.imwrite(output_path, image)
        # 保存裁剪后的图像

print("裁剪完成！")