{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":108394,"databundleVersionId":14167939,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport cv2\nimport glob\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\nfrom tensorflow.keras import layers, Model, mixed_precision, applications\nimport tensorflow.keras.backend as K\nfrom scipy.spatial.distance import pdist\nfrom sklearn.model_selection import train_test_split\nimport random\nimport itertools\n\n# --- 1. 全局配置 ---\nCONFIG = {\n    \"IMG_SIZE\": (128, 128),    # 图像尺寸，128足够看清纹理，太大爆显存\n    \"NUM_POINTS\": 1024,        # 点云采样数\n    \"BATCH_SIZE\": 32,          # 批次大小\n    \"EMBEDDING_DIM\": 128,      # 特征向量维度\n    \"EPOCHS\": 15,              # 训练轮数\n    \"LR\": 1e-4,                # 学习率\n    \"BASE_PATH\": \"/kaggle/input/h690/h690/h690\", # 根据实际路径修改\n    \"SEED\": 42\n}\n\n# --- 2. 开启混合精度训练 (Mixed Precision) ---\ntry:\n    policy = mixed_precision.Policy('mixed_float16')\n    mixed_precision.set_global_policy(policy)\n    print(\"🚀 混合精度训练已开启 (FP16) - 速度提升中...\")\nexcept:\n    print(\"⚠️ 无法开启混合精度，将使用默认 FP32\")\n\n# 固定随机种子\nnp.random.seed(CONFIG[\"SEED\"])\ntf.random.set_seed(CONFIG[\"SEED\"])","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_and_normalize_contour(image_path, target_points=1024):\n    \"\"\"\n    读取图像 -> 提取轮廓 -> 采样为固定点数 -> 归一化 -> 添加Z轴\n    \"\"\"\n    try:\n        img = cv2.imread(image_path)\n        if img is None: return np.zeros((target_points, 3))\n\n        # 1. 预处理与边缘检测\n        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n        blurred = cv2.GaussianBlur(gray, (5, 5), 0)\n        edges = cv2.Canny(blurred, 50, 150)\n        \n        # 闭运算连接断裂边缘\n        kernel = np.ones((5,5), np.uint8)\n        edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)\n\n        # 2. 提取轮廓\n        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n        if not contours: return np.zeros((target_points, 3))\n        \n        # 取最大的轮廓\n        c = max(contours, key=cv2.contourArea)\n        points = c.squeeze()\n        if points.ndim == 1: points = np.expand_dims(points, axis=0) # 处理单点情况\n\n        # 3. 采样/重采样到固定点数\n        num_points = len(points)\n        if num_points == 0: return np.zeros((target_points, 3))\n        \n        if num_points >= target_points:\n            indices = np.linspace(0, num_points - 1, target_points, dtype=int)\n            sampled = points[indices]\n        else:\n            # 点不够，重复填充\n            choice = np.random.choice(num_points, target_points, replace=True)\n            sampled = points[choice]\n\n        # 4. 归一化 (从像素坐标转换到 -1~1 空间)\n        centroid = np.mean(sampled, axis=0)\n        sampled = sampled - centroid # 中心化\n        m = np.max(np.sqrt(np.sum(sampled**2, axis=1))) # 最大距离\n        if m > 0: sampled = sampled / m\n            \n        # 5. 转 3D (添加 Z 轴，模拟厚度)\n        z = np.zeros((target_points, 1)) # 假设是一个平面切片\n        point_cloud = np.hstack([sampled, z])\n        \n        return point_cloud.astype(np.float32)\n\n    except Exception as e:\n        # print(f\"Error processing {image_path}: {e}\")\n        return np.zeros((target_points, 3), dtype=np.float32)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. 扫描所有图片\nIMAGE_DIR = os.path.join(CONFIG[\"BASE_PATH\"], 'sherd_images')\nall_image_paths = glob.glob(os.path.join(IMAGE_DIR, '*.jpg'))\nprint(f\"🔍 发现图片数量: {len(all_image_paths)}\")\n\n# 2. 预计算所有点云 (这步可能需要几分钟，但训练时会很快)\nprint(\"⏳ 正在预计算几何点云特征...\")\nPOINT_CLOUD_CACHE = {}\nfor path in all_image_paths: # 如果内存爆了，可以用 [:5000] 先测试\n    sherd_id = os.path.basename(path).split('.')[0] # 例如 JD0001_exterior\n    pc = extract_and_normalize_contour(path, CONFIG[\"NUM_POINTS\"])\n    POINT_CLOUD_CACHE[sherd_id] = pc\nprint(\"✅ 点云预处理完成！\")\n\n# 3. 生成训练对 (Demo 逻辑，你需要根据真实的 CSV 修改)\n# 假设 CSV 有 columns: [image_id, group_id]\n# group_id 相同的为正样本，不同的为负样本\ndef generate_pairs(metadata_df):\n    # 这里模拟从 metadata_df 读取\n    # 实际请用 pd.read_csv 加载你的 jd_sherds_info.csv\n    # 下面是伪代码逻辑，请确保你的 DataFrame 已经加载\n    \n    positive_pairs = []\n    negative_pairs = []\n    \n    # 示例：从缓存的 ID 中构建列表\n    available_ids = list(POINT_CLOUD_CACHE.keys())\n    \n    # --- 临时模拟数据 (请替换为真实 Group 逻辑) ---\n    # 真实逻辑：\n    # groups = df.groupby('unit')['image_id'].apply(list)\n    # for g in groups: \n    #    positive_pairs.extend(itertools.combinations(g, 2))\n    \n    # 这里为了代码能跑，我随机造一些对\n    print(\"⚠️ 正在生成训练样本对 (请替换为真实 CSV 逻辑)...\")\n    import random\n    for _ in range(5000): # 5000对正样本\n        id1 = random.choice(available_ids)\n        id2 = id1 # 简单模拟：自己跟自己是正样本 (实际应该是同组碎片)\n        positive_pairs.append((id1, id2, 1.0))\n        \n    for _ in range(5000): # 5000对负样本\n        id1 = random.choice(available_ids)\n        id2 = random.choice(available_ids)\n        if id1 != id2:\n            negative_pairs.append((id1, id2, 0.0))\n            \n    all_pairs = positive_pairs + negative_pairs\n    np.random.shuffle(all_pairs)\n    return all_pairs\n\n# 请确保这里加载了你的真实 csv\n# df = pd.read_csv(...)\n# training_pairs = generate_pairs(df)\ntraining_pairs = generate_pairs(None) # 使用模拟数据测试\nprint(f\"📊 生成训练对数量: {len(training_pairs)}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 定义层 ---\nclass L2Normalization(layers.Layer):\n    def call(self, x):\n        return tf.math.l2_normalize(x, axis=1)\n\nclass L2Distance(layers.Layer):\n    def call(self, inputs):\n        x, y = inputs\n        return K.sqrt(K.sum(K.square(x - y), axis=1, keepdims=True) + K.epsilon())\n\ndef mlp_block(x, filters):\n    x = layers.Conv1D(filters, 1)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Activation('relu')(x)\n    return x\n\n# --- T-Net (输入对齐) ---\ndef t_net(inputs, num_features):\n    x = mlp_block(inputs, 64)\n    x = mlp_block(x, 128)\n    x = mlp_block(x, 1024)\n    x = layers.GlobalMaxPooling1D()(x)\n    x = layers.Dense(512, activation='relu')(x)\n    x = layers.Dense(256, activation='relu')(x)\n    # 初始化为单位矩阵\n    bias = tf.keras.initializers.Constant(np.eye(num_features).flatten())\n    x = layers.Dense(num_features*num_features, kernel_initializer='zeros', bias_initializer=bias)(x)\n    x = layers.Reshape((num_features, num_features))(x)\n    return x\n\n# --- 1. 几何编码器 (PointNet) ---\ndef build_geometry_encoder():\n    inp = layers.Input(shape=(CONFIG[\"NUM_POINTS\"], 3))\n    \n    # T-Net 旋转对齐\n    transform = t_net(inp, 3)\n    x = layers.Dot(axes=(2, 1))([inp, transform])\n    \n    # 特征提取\n    x = mlp_block(x, 64)\n    x = mlp_block(x, 64)\n    x = mlp_block(x, 128)\n    x = mlp_block(x, 1024)\n    \n    # 全局特征\n    x = layers.GlobalMaxPooling1D()(x)\n    \n    # 降维\n    x = layers.Dense(512, activation='relu')(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Dropout(0.3)(x)\n    x = layers.Dense(256, activation='relu')(x)\n    \n    return Model(inp, x, name=\"Geometry_Encoder\")\n\n# --- 2. 视觉编码器 (EfficientNet) ---\ndef build_visual_encoder():\n    inp = layers.Input(shape=(*CONFIG[\"IMG_SIZE\"], 3))\n    \n    # 使用 EfficientNetB0，去头\n    base = applications.EfficientNetB0(include_top=False, weights='imagenet', input_tensor=inp)\n    # 稍微解冻顶层进行微调\n    base.trainable = True \n    for layer in base.layers[:-20]: layer.trainable = False\n        \n    x = layers.GlobalAveragePooling2D()(base.output)\n    x = layers.Dense(256, activation='relu')(x)\n    x = layers.BatchNormalization()(x)\n    \n    return Model(inp, x, name=\"Visual_Encoder\")\n\n# --- 3. 融合孪生网络 ---\ndef build_fusion_siamese():\n    # 输入定义\n    img_a = layers.Input(shape=(*CONFIG[\"IMG_SIZE\"], 3), name=\"img_a\")\n    pc_a  = layers.Input(shape=(CONFIG[\"NUM_POINTS\"], 3), name=\"pc_a\")\n    \n    img_b = layers.Input(shape=(*CONFIG[\"IMG_SIZE\"], 3), name=\"img_b\")\n    pc_b  = layers.Input(shape=(CONFIG[\"NUM_POINTS\"], 3), name=\"pc_b\")\n    \n    # 实例化编码器\n    vis_enc = build_visual_encoder()\n    geo_enc = build_geometry_encoder()\n    \n    # 提取特征 (共享权重)\n    feat_v_a = vis_enc(img_a)\n    feat_g_a = geo_enc(pc_a)\n    \n    feat_v_b = vis_enc(img_b)\n    feat_g_b = geo_enc(pc_b)\n    \n    # 特征融合 (Concatenate)\n    # 分支 A\n    x_a = layers.Concatenate()([feat_v_a, feat_g_a])\n    x_a = layers.Dense(256, activation='relu')(x_a)\n    x_a = layers.Dense(CONFIG[\"EMBEDDING_DIM\"])(x_a) # 确保它是 float32\n    emb_a = L2Normalization()(x_a)\n    \n    # 分支 B\n    x_b = layers.Concatenate()([feat_v_b, feat_g_b])\n    x_b = layers.Dense(256, activation='relu')(x_b)\n    x_b = layers.Dense(CONFIG[\"EMBEDDING_DIM\"])(x_b)\n    emb_b = L2Normalization()(x_b)\n    \n    # 计算距离\n    dist = L2Distance()([emb_a, emb_b])\n    \n    return Model(inputs=[img_a, pc_a, img_b, pc_b], outputs=dist, name=\"Fusion_Siamese\")\n\n# --- 4. 损失函数 (Contrastive Loss) ---\ndef contrastive_loss(y_true, y_pred):\n    margin = 1.0\n    y_true = tf.cast(y_true, tf.float32)\n    # y=1 (同类): minimize distance\n    # y=0 (异类): maximize distance (until margin)\n    sq_pred = K.square(y_pred)\n    margin_sq = K.square(K.maximum(margin - y_pred, 0))\n    return K.mean(y_true * sq_pred + (1 - y_true) * margin_sq)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class PotteryGenerator(tf.keras.utils.Sequence):\n    def __init__(self, pairs, pc_cache, image_dir, batch_size=32, img_size=(128,128), shuffle=True):\n        self.pairs = pairs\n        self.pc_cache = pc_cache\n        self.image_dir = image_dir\n        self.batch_size = batch_size\n        self.img_size = img_size\n        self.shuffle = shuffle\n        self.indexes = np.arange(len(self.pairs))\n        self.on_epoch_end()\n\n    def __len__(self):\n        return int(np.floor(len(self.pairs) / self.batch_size))\n\n    def on_epoch_end(self):\n        if self.shuffle:\n            np.random.shuffle(self.indexes)\n\n    def _load_img(self, sherd_id):\n        # 根据ID找图片，这里假设了命名规则，如果你的ID不带后缀，需要适配\n        # 优先找 exterior\n        path = os.path.join(self.image_dir, f\"{sherd_id}.jpg\")\n        if not os.path.exists(path):\n            # 尝试查找带后缀的\n            candidates = glob.glob(os.path.join(self.image_dir, f\"{sherd_id}*.jpg\"))\n            path = candidates[0] if candidates else None\n            \n        if path:\n            img = cv2.imread(path)\n            img = cv2.resize(img, self.img_size)\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            return applications.efficientnet.preprocess_input(img)\n        else:\n            # 找不到图片返回全黑\n            return np.zeros((*self.img_size, 3))\n\n    def __getitem__(self, index):\n        idxs = self.indexes[index*self.batch_size : (index+1)*self.batch_size]\n        batch_pairs = [self.pairs[k] for k in idxs]\n        \n        IA, PC_A, IB, PC_B = [], [], [], []\n        Labels = []\n        \n        for id1, id2, lbl in batch_pairs:\n            IA.append(self._load_img(id1))\n            IB.append(self._load_img(id2))\n            \n            # 从缓存获取点云，如果没有则返回零矩阵\n            PC_A.append(self.pc_cache.get(id1, np.zeros((CONFIG[\"NUM_POINTS\"], 3))))\n            PC_B.append(self.pc_cache.get(id2, np.zeros((CONFIG[\"NUM_POINTS\"], 3))))\n            \n            Labels.append(lbl)\n            \n        return [np.array(IA), np.array(PC_A), np.array(IB), np.array(PC_B)], np.array(Labels)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. 构建模型\nmodel = build_fusion_siamese()\nmodel.compile(\n    optimizer=tf.keras.optimizers.Adam(CONFIG[\"LR\"]),\n    loss=contrastive_loss\n)\n# model.summary() # 很大，可以不打印\n\n# 2. 划分训练集/验证集\ntrain_pairs, val_pairs = train_test_split(training_pairs, test_size=0.2, random_state=CONFIG[\"SEED\"])\n\n# 3. 实例化生成器\ntrain_gen = PotteryGenerator(train_pairs, POINT_CLOUD_CACHE, IMAGE_DIR, CONFIG[\"BATCH_SIZE\"])\nval_gen = PotteryGenerator(val_pairs, POINT_CLOUD_CACHE, IMAGE_DIR, CONFIG[\"BATCH_SIZE\"], shuffle=False)\n\n# 4. Callbacks (保存最佳模型，动态调整学习率)\ncallbacks = [\n    tf.keras.callbacks.ModelCheckpoint(\"best_pottery_model.h5\", save_best_only=True, monitor=\"val_loss\"),\n    tf.keras.callbacks.ReduceLROnPlateau(monitor=\"val_loss\", factor=0.5, patience=3, verbose=1),\n    tf.keras.callbacks.EarlyStopping(monitor=\"val_loss\", patience=6, restore_best_weights=True)\n]\n\nprint(\"🔥 开始训练多模态融合网络...\")\nhistory = model.fit(\n    train_gen,\n    validation_data=val_gen,\n    epochs=CONFIG[\"EPOCHS\"],\n    callbacks=callbacks,\n    verbose=1\n)\n\n# 5. 绘制结果\nplt.plot(history.history['loss'], label='Train Loss')\nplt.plot(history.history['val_loss'], label='Val Loss')\nplt.title('Training and Validation Loss')\nplt.legend()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}