import cv2
import numpy as np
import os
def fit_ellipse_to_mask(mask, min_area=1000):
    """对掩码进行椭圆拟合，返回拟合后的掩码"""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours: return mask
    largest_contour = max(contours, key=lambda c: cv2.contourArea(c))
    if cv2.contourArea(largest_contour) < min_area: return mask
    if len(largest_contour) >= 5:
        ellipse = cv2.fitEllipse(largest_contour)
        fitted_mask = np.zeros_like(mask)
        cv2.ellipse(fitted_mask, ellipse, 255, -1)
        return fitted_mask
    else:
        return mask
def contour_to_mask(contour, shape, fill_value=255):
    """
    将轮廓内所有点归为掩码
    
    参数:
        contour: numpy数组，形状为(N, 1, 2)，表示轮廓点集
        shape: 元组 (height, width)，指定输出掩码的大小
        fill_value: 掩码填充值，默认为255（白色）
        
    返回:
        numpy数组，二值掩码图像
    """
    # 创建一个全0的单通道图像（黑色背景）
    mask = np.zeros(shape, dtype=np.uint8)
    
    # 在掩码上绘制轮廓并填充内部区域
    cv2.drawContours(mask, [contour], -1, fill_value, -1)
    
    return mask
def filter_small_regions(mask, min_area=1000):
    """过滤掩码中面积小于指定阈值的区域"""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    filtered_mask = np.zeros_like(mask)
    
    for contour in contours:
        area = cv2.contourArea(contour)
        if area >= min_area:
            cv2.drawContours(filtered_mask, [contour], -1, 255, -1)
    
    return filtered_mask

def keep_center_region(mask, min_area=1800):
    """保留掩码中最靠近中心的区域"""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours: return mask
    h, w = mask.shape
    img_center = (w//2, h//2)
    best_contour = None
    min_distance = float('inf')
    for cnt in contours:
        if cv2.contourArea(cnt) < min_area: continue
        M = cv2.moments(cnt)
        if M["m00"] == 0: continue
        cx, cy = int(M["m10"]/M["m00"]), int(M["m01"]/M["m00"])
        distance = np.sqrt((cx-img_center[0])**2 + (cy-img_center[1])**2)
        if distance < min_distance:
            min_distance = distance
            best_contour = cnt
    if best_contour is None: return mask
    filtered_mask = np.zeros_like(mask)
    cv2.drawContours(filtered_mask, [best_contour], -1, 255, -1)
    return filtered_mask

def detect_egg_region(frame, min_area=10000, max_area=100000, visualize=True):
    """选择最大轮廓并进行椭圆拟合，将椭圆外像素置为30并可视化"""
    # 图像预处理（保持原有逻辑不变）
    lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    cl = clahe.apply(l)
    limg = cv2.merge((cl,a,b))
    enhanced_frame = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
    
    gray = cv2.cvtColor(enhanced_frame, cv2.COLOR_BGR2GRAY)
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                 cv2.THRESH_BINARY_INV, 11, 2)
    blur = cv2.GaussianBlur(thresh, (5, 5), 0)
    edges = cv2.Canny(blur, 20, 80)
    
    # 形态学操作闭合边缘
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    closed_edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
    
    # 寻找轮廓并筛选最大轮廓
    contours, _ = cv2.findContours(closed_edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return np.full_like(frame, 37), None, edges  # 返回全30图像
    
    egg_contour = max(contours, key=lambda c: cv2.contourArea(c))
    area = cv2.contourArea(egg_contour)
    
    if not (min_area < area < max_area):
        return np.full_like(frame, 37), None, edges  # 返回全30图像
    
    # ------------------- 新增椭圆拟合逻辑 -------------------
    # 对最大轮廓进行椭圆拟合（需至少5个顶点）
    if len(egg_contour) >= 5:
        ellipse = cv2.fitEllipse(egg_contour)  # 拟合椭圆
        # 创建椭圆掩码
        egg_mask = np.zeros_like(gray)
        cv2.ellipse(egg_mask, ellipse, 255, -1)  # -1 表示填充椭圆内部
    else:
        # 轮廓顶点不足时，使用原轮廓创建掩码
        egg_mask = np.zeros_like(gray)
        cv2.drawContours(egg_mask, [egg_contour], -1, 255, -1)
    
    # 形态学闭操作优化掩码（可选，根据需要调整核大小）
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
    egg_mask = cv2.morphologyEx(egg_mask, cv2.MORPH_CLOSE, kernel)
    # ------------------------------------------------------
    
    # 将椭圆外的像素置为30
    masked_frame = frame.copy()
    masked_frame[egg_mask == 0] = 37
    
    # 可视化：显示椭圆拟合结果
    if visualize:
        cv2.imshow("Egg Ellipse Mask", egg_mask)  # 显示椭圆掩码
        cv2.imshow("Masked Frame with Ellipse", masked_frame)
    
    return masked_frame

def process_frame(frame, 
                  background_thresh=40,  
                  egg_white_thresh=120,   
                  yolk_thresh=200,      
                  germ_thresh=180,      
                  alpha=0.2,
                  visualize=True):  # 添加visualize参数
    """使用固定阈值处理有蛋清的图像"""
    # 转换为灰度图
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    
    # 预处理：高斯模糊减少噪点
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # 增强对比度
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    gray = clahe.apply(gray)
    
    # 创建背景掩码（最暗区域）
    mask_background = cv2.threshold(gray, background_thresh, 255, cv2.THRESH_BINARY_INV)[1]
    
    # 固定阈值分割（严格按照：背景 < 蛋清 < 蛋黄 < 胚盘）
    mask_egg_white = cv2.inRange(gray, background_thresh + 1, egg_white_thresh)
    mask_yolk = cv2.inRange(gray, egg_white_thresh + 1, yolk_thresh)
    mask_germ = cv2.threshold(gray, germ_thresh, 255, cv2.THRESH_BINARY)[1]
    mask_germ=filter_small_regions(mask_germ)
    # 形态学优化
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    mask_egg_white = cv2.morphologyEx(mask_egg_white, cv2.MORPH_CLOSE, kernel, iterations=2)
    mask_yolk = cv2.morphologyEx(mask_yolk, cv2.MORPH_CLOSE, kernel)
    #mask_germ = cv2.morphologyEx(mask_germ, cv2.MORPH_CLOSE, kernel)
    cv2.imshow("mask_germ", mask_germ)
    # 确保掩码之间没有重叠区域
    mask_yolk = cv2.bitwise_and(mask_yolk, cv2.bitwise_not(mask_egg_white))
    mask_yolk = cv2.bitwise_and(mask_yolk, cv2.bitwise_not(mask_germ))
    #mask_germ = cv2.bitwise_and(mask_germ, cv2.bitwise_not(mask_yolk))
    
    #mask_germ = cv2.bitwise_and(mask_germ, cv2.bitwise_not(mask_egg_white))
    
    # 新增：只保留蛋清中心区域
    #mask_egg_white_center = keep_center_region(mask_egg_white, min_area=5000)
    mask_egg_white_center = mask_egg_white
    # 椭圆拟合（使用处理后的蛋清掩码）
    mask_egg_white_fit = fit_ellipse_to_mask(mask_egg_white_center, min_area=10000)  # 增大最小面积
    mask_yolk_fit = fit_ellipse_to_mask(mask_yolk, min_area=4000)
    # 胚盘处理
    #mask_germ_filtered = keep_center_region(mask_germ, min_area=300)
    mask_germ_filtered =mask_germ
    
    # 创建彩色掩码
    color_mask = np.zeros_like(frame)
    color_mask[mask_egg_white_fit > 0] = (0, 0, 255)  # 红色蛋清
    color_mask[mask_yolk_fit > 0] = (0, 255, 0)      # 绿色蛋黄
    color_mask[mask_germ_filtered > 0] = (0, 255, 255)  # 黄色胚盘
    
    # 确保背景区域完全为黑色（最后应用背景掩码）
    color_mask[mask_background > 0] = (0, 0, 0)
    
    # 应用透明度
    result = cv2.addWeighted(frame, 1 - alpha, color_mask, alpha, 0)
    
    # 可视化
    if visualize:
        cv2.imshow("Processed Frame", result)
    
    return result

def main():
    input_path = "no\W-7.gif"
    cap = cv2.VideoCapture(input_path)
    
    if not cap.isOpened():
        print(f"Error: 无法打开文件 {input_path}")
        return
    
    # 获取原始文件名和扩展名
    base_name = os.path.basename(input_path)
    file_name, file_ext = os.path.splitext(base_name)
    
    # 构建输出文件名
    output_name = f"processed_{file_name}.mp4"
    
    fps = cap.get(cv2.CAP_PROP_FPS) or 10
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    out = cv2.VideoWriter(output_name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    
    frame_count = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        cv2.imshow("Original Frame", frame)
        new_frame = detect_egg_region(frame, min_area=10000, max_area=100000, visualize=True)
        
        processed_frame = process_frame(new_frame, 
                                        background_thresh=40,
                                        egg_white_thresh=120,
                                        yolk_thresh=200,
                                        germ_thresh=190,
                                        alpha=0.2,
                                        visualize=True)
        
        out.write(processed_frame)
        
        if cv2.waitKey(1) & 0xFF == 27:
            break
        
        frame_count += 1
        if frame_count % 10 == 0:
            print(f"已处理 {frame_count} 帧")
    
    cap.release()
    out.release()
    cv2.destroyAllWindows()
    print(f"处理完成，输出文件: {output_name}")

if __name__ == "__main__":
    main()
