"""
二維碼識(shí)別與傾斜糾正程序
功能：
1. 檢測圖像中的二維碼
2. 糾正傾斜姿態(tài)
3. 識(shí)別二維碼內(nèi)容
"""

import cv2
import numpy as np
import pyzbar.pyzbar as pyzbar
from pyzbar.pyzbar import ZBarSymbol


class QRCodeDetector:
    """二維碼檢測與糾正類"""
    
    def __init__(self):
        self.debug = False
    
    def detect_qr_corners(self, image):
        """
        檢測二維碼的四個(gè)角點(diǎn)
        
        Args:
            image: 輸入圖像 (BGR格式)
            
        Returns:
            corners_list: 檢測到的角點(diǎn)列表，每個(gè)元素是4個(gè)角點(diǎn)的數(shù)組
        """
        # 轉(zhuǎn)換為灰度圖
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用QRCodeDetector檢測二維碼
        qr_detector = cv2.QRCodeDetector()
        
        # 檢測并解碼
        retval, decoded_info, points, straight_qrcode = qr_detector.detectAndDecodeMulti(gray)
        
        corners_list = []
        if retval and points is not None:
            for i, pts in enumerate(points):
                if pts is not None:
                    corners = pts.reshape(4, 2).astype(np.float32)
                    corners_list.append(corners)
        
        return corners_list
    
    def detect_qr_by_contours(self, image):
        """
        使用輪廓檢測方法找到二維碼區(qū)域
        
        Args:
            image: 輸入圖像
            
        Returns:
            corners_list: 檢測到的角點(diǎn)列表
        """
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 自適應(yīng)閾值處理
        thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                       cv2.THRESH_BINARY, 51, 10)
        
        # 形態(tài)學(xué)操作
        kernel = np.ones(5, np.uint8)
        morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        
        # 查找輪廓
        contours, _ = cv2.findContours(morph, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        
        corners_list = []
        
        for contour in contours:
            # 計(jì)算輪廓面積
            area = cv2.contourArea(contour)
            if area < 1000:  # 過濾太小的區(qū)域
                continue
            
            # 多邊形逼近
            peri = cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
            
            # 尋找四邊形
            if len(approx) == 4:
                corners = approx.reshape(4, 2).astype(np.float32)
                
                # 檢查是否是正方形/矩形（二維碼形狀）
                angles = self._calculate_angles(corners)
                if all(70 < angle < 110 for angle in angles):
                    corners_list.append(self._order_points(corners))
        
        return corners_list
    
    def _calculate_angles(self, corners):
        """計(jì)算四邊形的內(nèi)角"""
        angles = []
        n = len(corners)
        for i in range(n):
            p1 = corners[(i - 1) % n]
            p2 = corners[i]
            p3 = corners[(i + 1) % n]
            
            v1 = p1 - p2
            v2 = p3 - p2
            
            angle = np.arctan2(v2[1], v2[0]) - np.arctan2(v1[1], v1[0])
            angle = np.degrees(angle)
            if angle < 0:
                angle += 360
            if angle > 180:
                angle = 360 - angle
            angles.append(angle)
        return angles
    
    def _order_points(self, pts):
        """
        對(duì)四個(gè)角點(diǎn)進(jìn)行排序：左上、右上、右下、左下
        """
        rect = np.zeros((4, 2), dtype=np.float32)
        
        # 計(jì)算每個(gè)點(diǎn)的坐標(biāo)和
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]  # 左上
        rect[2] = pts[np.argmax(s)]  # 右下
        
        # 計(jì)算差值
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]  # 右上
        rect[3] = pts[np.argmax(diff)]  # 左下
        
        return rect
    
    def correct_perspective(self, image, corners, output_size=None):
        """
        糾正透 視變換，將傾斜的二維碼轉(zhuǎn)正
        
        Args:
            image: 輸入圖像
            corners: 四個(gè)角點(diǎn) [左上, 右上, 右下, 左下]
            output_size: 輸出圖像大小 (寬, 高)，默認(rèn)自動(dòng)計(jì)算
            
        Returns:
            corrected_image: 糾正后的圖像
        """
        # 確保角點(diǎn)順序正確
        ordered_corners = self._order_points(corners)
        
        # 計(jì)算輸出尺寸
        if output_size is None:
            # 根據(jù)角點(diǎn)計(jì)算寬度和高度
            width_top = np.linalg.norm(ordered_corners[1] - ordered_corners[0])
            width_bottom = np.linalg.norm(ordered_corners[2] - ordered_corners[3])
            width = int(max(width_top, width_bottom))
            
            height_left = np.linalg.norm(ordered_corners[3] - ordered_corners[0])
            height_right = np.linalg.norm(ordered_corners[2] - ordered_corners[1])
            height = int(max(height_left, height_right))
            
            output_size = (width, height)
        
        # 定義目標(biāo)點(diǎn)
        dst_points = np.array([
            [0, 0],
            [output_size[0] - 1, 0],
            [output_size[0] - 1, output_size[1] - 1],
            [0, output_size[1] - 1]
        ], dtype=np.float32)
        
        # 計(jì)算透視變換矩陣
        matrix = cv2.getPerspectiveTransform(ordered_corners, dst_points)
        
        # 應(yīng)用透 視變換
        corrected = cv2.warpPerspective(image, matrix, output_size)
        
        return corrected
    
    def decode_qr(self, image):
        """
        解碼二維碼
        
        Args:
            image: 輸入圖像
            
        Returns:
            results: 解碼結(jié)果列表
        """
        # 轉(zhuǎn)換為灰度圖
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
        
        # 解碼二維碼
        decoded_objects = pyzbar.decode(gray, symbols=[ZBarSymbol.QRCODE])
        
        results = []
        for obj in decoded_objects:
            result = {
                'data': obj.data.decode('utf-8'),
                'type': obj.type,
                'rect': obj.rect,
                'polygon': obj.polygon
            }
            results.append(result)
        
        return results
    
    def process_image(self, image_path, save_corrected=False, output_path=None):
        """
        處理圖像：檢測、糾正、識(shí)別二維碼
        
        Args:
            image_path: 輸入圖像路徑
            save_corrected: 是否保存糾正后的圖像
            output_path: 糾正后圖像的保存路徑
            
        Returns:
            results: 識(shí)別結(jié)果列表
        """
        # 讀取圖像
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"無法讀取圖像: {image_path}")
        
        print(f"處理圖像: {image_path}")
        print(f"圖像尺寸: {image.shape[1]} x {image.shape[0]}")
        
        results = []
        
        # 首先嘗試直接識(shí)別（不糾正）
        direct_results = self.decode_qr(image)
        if direct_results:
            print("? 直接識(shí)別成功！")
            for i, result in enumerate(direct_results):
                print(f"  [{i+1}] 內(nèi)容: {result['data']}")
                results.append({
                    'method': 'direct',
                    'corrected': False,
                    **result
                })
            return results
        
        print("直接識(shí)別失敗，嘗試檢測并糾正傾斜...")
        
        # 嘗試檢測二維碼角點(diǎn)
        corners_list = self.detect_qr_corners(image)
        
        if not corners_list:
            print("使用備用方法檢測...")
            corners_list = self.detect_qr_by_contours(image)
        
        if not corners_list:
            print("? 未檢測到二維碼")
            return results
        
        print(f"檢測到 {len(corners_list)} 個(gè)可能的二維碼區(qū)域")
        
        # 對(duì)每個(gè)檢測到的區(qū)域進(jìn)行糾正和識(shí)別
        for i, corners in enumerate(corners_list):
            print(f"\n處理區(qū)域 {i+1}...")
            
            # 糾正透 視
            corrected_image = self.correct_perspective(image, corners)
            
            # 保存糾正后的圖像（如果需要）
            if save_corrected and output_path:
                corrected_path = output_path.replace('.png', f'_corrected_{i+1}.png')
                cv2.imwrite(corrected_path, corrected_image)
                print(f"  已保存糾正后的圖像: {corrected_path}")
            
            # 嘗試解碼糾正后的圖像
            decoded_results = self.decode_qr(corrected_image)
            
            if decoded_results:
                print(f"  ? 糾正后識(shí)別成功！")
                for j, result in enumerate(decoded_results):
                    print(f"    [{j+1}] 內(nèi)容: {result['data']}")
                    results.append({
                        'method': 'corrected',
                        'corrected': True,
                        'region_index': i + 1,
                        **result
                    })
            else:
                print(f"  ? 糾正后仍無法識(shí)別")
                
                # 嘗試對(duì)糾正后的圖像進(jìn)行增強(qiáng)
                enhanced = self._enhance_qr_image(corrected_image)
                decoded_results = self.decode_qr(enhanced)
                
                if decoded_results:
                    print(f"  ? 增強(qiáng)后識(shí)別成功！")
                    for j, result in enumerate(decoded_results):
                        print(f"    [{j+1}] 內(nèi)容: {result['data']}")
                        results.append({
                            'method': 'enhanced',
                            'corrected': True,
                            'region_index': i + 1,
                            **result
                        })
        
        return results
    
    def _enhance_qr_image(self, image):
        """
        增強(qiáng)二維碼圖像以提高識(shí)別率
        """
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
        
        # 嘗試多種預(yù)處理方法
        enhanced_images = []
        
        # 1. 自適應(yīng)閾值
        enhanced_images.append(cv2.adaptiveThreshold(gray, 255, 
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 21, 5))
        
        # 2. Otsu閾值
        _, otsu = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        enhanced_images.append(otsu)
        
        # 3. 直方圖均衡化
        enhanced_images.append(cv2.equalizeHist(gray))
        
        # 4. 高斯模糊后銳化
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        enhanced_images.append(cv2.addWeighted(gray, 1.5, blurred, -0.5, 0))
        
        # 嘗試每種增強(qiáng)方法
        for enhanced in enhanced_images:
            results = pyzbar.decode(enhanced, symbols=[ZBarSymbol.QRCODE])
            if results:
                return enhanced
        
        return gray


def main():
    """主函數(shù)"""
    import argparse
    import os
    
    parser = argparse.ArgumentParser(description='二維碼識(shí)別與傾斜糾正工具')
    parser.add_argument('image', help='輸入圖像路徑')
    parser.add_argument('-s', '--save', action='store_true', help='保存糾正后的圖像')
    parser.add_argument('-o', '--output', help='輸出圖像路徑')
    parser.add_argument('-d', '--debug', action='store_true', help='啟用調(diào)試模式')
    
    args = parser.parse_args()
    
    # 檢查輸入文件是否存在
    if not os.path.exists(args.image):
        print(f"錯(cuò)誤: 文件不存在 - {args.image}")
        return
    
    # 創(chuàng)建檢測器
    detector = QRCodeDetector()
    detector.debug = args.debug
    
    # 處理圖像
    try:
        results = detector.process_image(
            args.image, 
            save_corrected=args.save, 
            output_path=args.output
        )
        
        print("\n" + "="*50)
        print("識(shí)別結(jié)果匯總:")
        print("="*50)
        
        if results:
            for i, result in enumerate(results):
                print(f"\n[{i+1}] 識(shí)別方法: {result['method']}")
                print(f"    內(nèi)容: {result['data']}")
                print(f"    類型: {result['type']}")
                if result.get('corrected'):
                    print(f"    已糾正: 是")
        else:
            print("未能識(shí)別任何二維碼")
            
    except Exception as e:
        print(f"處理過程中出錯(cuò): {e}")
        import traceback
        traceback.print_exc()


if __name__ == '__main__':
    main()