import cv2
import numpy as np
import pyzbar.pyzbar as pyzbar
from torchvision import transforms, models
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from PIL import Image
import os

# 1. 傳統(tǒng)條形碼識(shí)別方法實(shí)現(xiàn)
def traditional_barcode_detection(image_path):
    """使用pyzbar實(shí)現(xiàn)傳統(tǒng)條形碼識(shí)別"""
    # 讀取圖像
    image = cv2.imread(image_path)
    if image is None:
        return "無法讀取圖像"
    
    # 轉(zhuǎn)為灰度圖
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # 使用pyzbar檢測(cè)條形碼
    barcodes = pyzbar.decode(gray)
    
    result = []
    for barcode in barcodes:
        # 提取條形碼位置
        (x, y, w, h) = barcode.rect
        cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
        
        # 解碼條形碼數(shù)據(jù)
        barcode_data = barcode.data.decode("utf-8")
        barcode_type = barcode.type
        
        # 顯示結(jié)果
        text = f"{barcode_type}: {barcode_data}"
        cv2.putText(image, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        result.append({
            "type": barcode_type,
            "data": barcode_data,
            "position": (x, y, w, h)
        })
    
    # 保存結(jié)果圖像
    cv2.imwrite("traditional_barcode_result.jpg", image)
    return result

# 2. 深度學(xué)習(xí)模型：基于CNN的條形碼檢測(cè)與分類
class BarcodeDetector(nn.Module):
    def __init__(self, num_classes=10):  # num_classes取決于條形碼類型數(shù)量
        super(BarcodeDetector, self).__init__()
        # 使用預(yù)訓(xùn)練的ResNet作為特征提取器
        self.feature_extractor = models.resnet18(pretrained=True)
        # 替換最后一層以適應(yīng)我們的分類任務(wù)
        num_ftrs = self.feature_extractor.fc.in_features
        self.feature_extractor.fc = nn.Linear(num_ftrs, num_classes)
        
        # 可以添加目標(biāo)檢測(cè)層如Faster R-CNN或YOLO，但這里簡(jiǎn)化為分類模型
        # 實(shí)際應(yīng)用中應(yīng)使用目標(biāo)檢測(cè)模型定位條形碼區(qū)域
    
    def forward(self, x):
        return self.feature_extractor(x)

# 3. 數(shù)據(jù)處理與增強(qiáng)
def preprocess_image(image_path, target_size=(224, 224)):
    """圖像預(yù)處理函數(shù)，包括縮放、歸一化等"""
    image = Image.open(image_path).convert('RGB')
    transform = transforms.Compose([
        transforms.Resize(target_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)  # 添加批次維度

def data_augmentation(image_dir, output_dir, num_augmentations=5):
    """數(shù)據(jù)增強(qiáng)函數(shù)，生成更多訓(xùn)練樣本"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    for filename in os.listdir(image_dir):
        if filename.endswith(('.jpg', '.jpeg', '.png')):
            image_path = os.path.join(image_dir, filename)
            image = Image.open(image_path).convert('RGB')
            
            # 定義數(shù)據(jù)增強(qiáng)操作
            augmentations = transforms.Compose([
                transforms.RandomRotation(10),
                transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
                transforms.RandomHorizontalFlip(),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2)
            ])
            
            # 生成增強(qiáng)樣本
            base_name, ext = os.path.splitext(filename)
            for i in range(num_augmentations):
                augmented = augmentations(image)
                output_path = os.path.join(output_dir, f"{base_name}_aug_{i}{ext}")
                transforms.ToPILImage()(augmented).save(output_path)

# 4. 模型訓(xùn)練函數(shù)
def train_model(model, train_loader, val_loader, num_epochs=10, learning_rate=0.001):
    """訓(xùn)練深度學(xué)習(xí)模型"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    
    for epoch in range(num_epochs):
        # 訓(xùn)練階段
        model.train()
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
        
        # 驗(yàn)證階段
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        print(f"Epoch {epoch+1}/{num_epochs}, "
              f"Train Loss: {running_loss/len(train_loader):.4f}, "
              f"Val Loss: {val_loss/len(val_loader):.4f}, "
              f"Val Accuracy: {100 * correct / total:.2f}%")
    
    return model

# 5. 結(jié)合傳統(tǒng)方法和深度學(xué)習(xí)的高級(jí)識(shí)別函數(shù)
def advanced_barcode_recognition(image_path, model=None, use_deep_learning=True):
    """高級(jí)條形碼識(shí)別，結(jié)合傳統(tǒng)方法和深度學(xué)習(xí)"""
    # 讀取圖像
    image = cv2.imread(image_path)
    if image is None:
        return "無法讀取圖像"
    
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    if use_deep_learning and model:
        # 使用深度學(xué)習(xí)模型進(jìn)行條形碼檢測(cè)(簡(jiǎn)化示例，實(shí)際應(yīng)使用目標(biāo)檢測(cè)模型)
        # 這里僅演示流程，實(shí)際需要先定位條形碼區(qū)域
        processed_img = preprocess_image(image_path)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        model.eval()
        
        with torch.no_grad():
            outputs = model(processed_img.to(device))
            # 根據(jù)輸出判斷是否存在條形碼及類型
            # 這里需要實(shí)際的分類邏輯
            _, predicted = torch.max(outputs.data, 1)
            # ... 更多處理 ...
            
    else:
        # 使用傳統(tǒng)方法
        barcodes = pyzbar.decode(gray)
        result = []
        for barcode in barcodes:
            (x, y, w, h) = barcode.rect
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
            
            barcode_data = barcode.data.decode("utf-8")
            barcode_type = barcode.type
            
            text = f"{barcode_type}: {barcode_data}"
            cv2.putText(image, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            
            result.append({
                "type": barcode_type,
                "data": barcode_data,
                "position": (x, y, w, h)
            })
        
        cv2.imwrite("advanced_barcode_result.jpg", image)
        return result
    
    return "示例中未實(shí)現(xiàn)完整的深度學(xué)習(xí)檢測(cè)流程，可參考傳統(tǒng)方法結(jié)果"

# 示例使用
if __name__ == "__main__":
    # 測(cè)試傳統(tǒng)方法
    print("使用傳統(tǒng)方法識(shí)別條形碼:")
    traditional_result = traditional_barcode_detection("barcode_test.jpg")
    if traditional_result:
        for barcode in traditional_result:
            print(f"類型: {barcode['type']}, 數(shù)據(jù): {barcode['data']}")
    else:
        print("未檢測(cè)到條形碼")
    
    # 注意：深度學(xué)習(xí)模型的訓(xùn)練和使用需要準(zhǔn)備數(shù)據(jù)集和完整實(shí)現(xiàn)
    # 這里僅展示框架和流程