import cv2
import numpy as np

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 keep_center_region(mask, min_area=500):
    """保留掩码中最靠近中心的区域"""
    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 process_frame_without_eggwhite(frame, 
                                  yolk_thresh=130, 
                                  germ_thresh=130,
                                  brightness_weight=0.8,
                                  alpha=0.5):
    """无蛋清状态处理函数（恢复原始逻辑）"""
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    mask_background = (gray == 0).astype(np.uint8) * 255
    
    # 动态阈值调整
    non_background_gray = gray[gray != 0]
    if non_background_gray.size == 0:
        avg_brightness = 0
    else:
        avg_brightness = np.mean(non_background_gray)
    
    brightness_ratio = avg_brightness / 128
    adjusted_yolk_thresh = int(yolk_thresh * brightness_ratio * brightness_weight)
    adjusted_germ_thresh = int(germ_thresh * brightness_ratio * brightness_weight)
    
    yolk_thresh = max(50, min(200, adjusted_yolk_thresh))
    germ_thresh = max(80, min(250, adjusted_germ_thresh))
    
    # 掩码生成（无蛋清，仅蛋黄和胚盘）
    mask_yolk = cv2.inRange(gray, 1, yolk_thresh)
    mask_germ = cv2.threshold(gray, germ_thresh, 255, cv2.THRESH_BINARY)[1]
    
    # 形态学优化
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask_yolk = cv2.morphologyEx(mask_yolk, cv2.MORPH_CLOSE, kernel)
    mask_germ = cv2.morphologyEx(mask_germ, cv2.MORPH_CLOSE, kernel)
    
    # 椭圆拟合（仅蛋黄）
    mask_yolk_fit = fit_ellipse_to_mask(mask_yolk, min_area=2000)
    
    # 胚盘处理
    mask_germ_filtered = keep_center_region(mask_germ, min_area=500)
    
    # 创建彩色掩码（无蛋清，无红色区域）
    color_mask = np.zeros_like(frame)
    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)
    
    return result

def process_frame_with_eggwhite(frame, 
                               egg_white_thresh=80, 
                               yolk_thresh=130, 
                               germ_thresh=130,
                               alpha=0.5):
    """有蛋清状态处理函数（保持原始逻辑）"""
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    mask_background = (gray == 0).astype(np.uint8) * 255
    
    # 固定阈值
    mask_egg_white = cv2.inRange(gray, 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]
    
    # 形态学优化
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask_egg_white = cv2.morphologyEx(mask_egg_white, cv2.MORPH_CLOSE, kernel)
    mask_yolk = cv2.morphologyEx(mask_yolk, cv2.MORPH_CLOSE, kernel)
    mask_germ = cv2.morphologyEx(mask_germ, cv2.MORPH_CLOSE, kernel)
    
    # 椭圆拟合
    mask_egg_white_fit = fit_ellipse_to_mask(mask_egg_white, min_area=5000)
    mask_yolk_fit = fit_ellipse_to_mask(mask_yolk, min_area=2000)
    
    # 胚盘处理
    mask_germ_filtered = keep_center_region(mask_germ, min_area=500)
    
    # 创建彩色掩码
    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)
    
    return result

def process_frame(frame, 
                  egg_white_thresh=65, 
                  yolk_thresh=128, 
                  germ_thresh=128,
                  brightness_weight=0.8,
                  alpha=0.3):
    """主处理函数：根据蛋清状态选择不同处理路径"""
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    
    # 仅使用原始固定阈值检测蛋清
    mask_egg_white = cv2.inRange(gray, 1, egg_white_thresh)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask_egg_white = cv2.morphologyEx(mask_egg_white, cv2.MORPH_CLOSE, kernel)
    
    # 检测蛋清存在
    egg_white_exists = np.sum(mask_egg_white) > 30000000  # 原始检测阈值
    
    # 根据状态选择处理函数
    if egg_white_exists:
        return process_frame_with_eggwhite(frame, egg_white_thresh, yolk_thresh, germ_thresh, alpha), True
    else:
        return process_frame_without_eggwhite(frame, yolk_thresh, germ_thresh, brightness_weight, alpha), False

def main():
    cap = cv2.VideoCapture("20250522-162702.mp4")
    if not cap.isOpened():
        print("Error: 无法打开视频文件")
        return
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    out = cv2.VideoWriter("output_colored.avi", cv2.VideoWriter_fourcc(*'XVID'), fps, (width, height))
    
    frame_count = 0
    alpha = 0.2  # 设置固定透明度值
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # 处理帧并获取蛋清状态
        processed_frame, egg_white_exists = process_frame(frame, alpha=alpha)
        
        out.write(processed_frame)
        cv2.imshow("Processed Frame", processed_frame)
        
        if cv2.waitKey(1) & 0xFF == 27:  # ESC退出
            break
        
        frame_count += 1
        if frame_count % 10 == 0:
            print(f"已处理 {frame_count} 帧，状态: {'有蛋清' if egg_white_exists else '无蛋清'}")
    
    cap.release()
    out.release()
    cv2.destroyAllWindows()
    print(f"处理完成，共处理 {frame_count} 帧")

if __name__ == "__main__":
    main()
