# 1. 安装依赖库（首次使用需执行）
# !pip install datasets pillow opencv-python numpy

# 2. 导入工具库
from datasets import load_dataset
import matplotlib.pyplot as plt
import cv2
import numpy as np

# 3. 加载数据集（二选一：完整数据集/调试子集）
# 加载完整修正单评者数据集（约59k细胞核，1744条目，需较多存储空间）
dataset = load_dataset("minhanhto09/NuCLS_dataset", name="default")
# 加载调试子集（仅train_fold_999/test_fold_999，28条目，快速测试）
# dataset = load_dataset("minhanhto09/NuCLS_dataset", name="debug")

# 4. 查看数据集结构（训练/测试划分、条目数量）
print("数据集拆分：", dataset.keys())  # 输出：dict_keys(['train_fold_1', 'test_fold_1', ..., 'test_fold_999'])
print("train_fold_1条目数：", len(dataset["train_fold_1"]))  # 输出：1481
print("单条数据字段：", dataset["train_fold_1"][0].keys())  # 输出：file_name, rgb_image, mask_image, ...

# 5. 可视化单条数据（RGB图+掩码图+叠加图）
sample = dataset["train_fold_1"][0]
fig, axes = plt.subplots(1, 3, figsize=(15, 5))

# RGB图（PIL格式转OpenCV格式便于显示）
rgb_img = np.array(sample["rgb_image"])
axes[0].imshow(rgb_img)
axes[0].set_title("RGB原图")
axes[0].axis("off")

# 掩码图（第一通道：类别标签；第二三通道：细胞核唯一ID）
mask_img = np.array(sample["mask_image"])
axes[1].imshow(mask_img[:, :, 0], cmap="tab10")  # 显示类别标签通道
axes[1].set_title("掩码图（类别通道）")
axes[1].axis("off")

# 可视化叠加图
vis_img = np.array(sample["visualization_image"])
axes[2].imshow(vis_img)
axes[2].set_title("RGB+掩码叠加图")
axes[2].axis("off")

plt.tight_layout()
plt.show()

# 6. 查看标注坐标与分类标签（以第一个细胞核为例）
annotations = sample["annotation_coordinates"]
print("细胞核标注数量：", len(annotations))  # 输出该条目下的细胞核总数
print("第一个细胞核标签：")
print(f"- 细粒度分类（raw）：{annotations[0]['raw_classification']}")
print(f"- 中粒度分类（main）：{annotations[0]['main_classification']}")
print(f"- 粗粒度分类（super）：{annotations[0]['super_classification']}")
print(f"- 边界框坐标（xmin, ymin, xmax, ymax）：{annotations[0]['xmin']}, {annotations[0]['ymin']}, {annotations[0]['xmax']}, {annotations[0]['ymax']}")

# 1. 安装YOLO库
# !pip install ultralytics

# 2. 数据格式转换（YOLO要求：图像+标签文件（每行为“类别ID x_center y_center w h”））
import os
from ultralytics.data.utils import autosplit

# 定义类别映射（super_classification的4类→0-3ID）
super_class_map = {
    "sTIL": 0,
    "nonTIL_stromal": 1,
    "tumor": 2,
    "other": 3  # 覆盖剩余super分类（如数据集定义的其他类别）
}

# 创建YOLO格式数据集目录
yolo_data_dir = "nucls_yolo_data"
os.makedirs(os.path.join(yolo_data_dir, "images", "train"), exist_ok=True)
os.makedirs(os.path.join(yolo_data_dir, "labels", "train"), exist_ok=True)
os.makedirs(os.path.join(yolo_data_dir, "images", "val"), exist_ok=True)
os.makedirs(os.path.join(yolo_data_dir, "labels", "val"), exist_ok=True)

# 处理train_fold_1数据（示例：取前1000条为训练，剩余为验证）
train_sample = dataset["train_fold_1"].select(range(1000))
val_sample = dataset["train_fold_1"].select(range(1000, len(dataset["train_fold_1"])))

def convert_to_yolo(sample, save_dir, split):
    for idx, data in enumerate(sample):
        # 保存RGB图像
        img = np.array(data["rgb_image"])
        img_path = os.path.join(save_dir, "images", split, f"nucls_{idx}.jpg")
        cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))  # OpenCV保存需BGR格式
        
        # 生成标签文件
        label_path = os.path.join(save_dir, "labels", split, f"nucls_{idx}.txt")
        with open(label_path, "w") as f:
            for ann in data["annotation_coordinates"]:
                # 获取粗粒度类别ID
                super_cls = ann["super_classification"]
                if super_cls not in super_class_map:
                    continue  # 跳过未定义类别
                cls_id = super_class_map[super_cls]
                
                # 图像尺寸（H, W）
                h, w = img.shape[0], img.shape[1]
                # 边界框坐标转换（xmin, ymin, xmax, ymax → x_center, y_center, w, h，归一化到0-1）
                x_center = (ann["xmin"] + ann["xmax"]) / (2 * w)
                y_center = (ann["ymin"] + ann["ymax"]) / (2 * h)
                bbox_w = (ann["xmax"] - ann["xmin"]) / w
                bbox_h = (ann["ymax"] - ann["ymin"]) / h
                
                # 写入标签（格式：cls_id x_center y_center w h）
                f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {bbox_w:.6f} {bbox_h:.6f}\n")

# 转换训练集与验证集
convert_to_yolo(train_sample, yolo_data_dir, "train")
convert_to_yolo(val_sample, yolo_data_dir, "val")

# 3. 编写YOLO配置文件（data.yaml）
yaml_content = f"""
path: {os.path.abspath(yolo_data_dir)}  # 数据集根路径
train: images/train  # 训练图像路径
val: images/val      # 验证图像路径

nc: 4  # 类别数量
names: ['sTIL', 'nonTIL_stromal', 'tumor', 'other']  # 类别名称
"""
with open(os.path.join(yolo_data_dir, "data.yaml"), "w") as f:
    f.write(yaml_content)

# 4. 训练YOLOv8n模型
from ultralytics import YOLO

# 加载预训练模型
model = YOLO("yolov8n.pt")

# 训练（重点参数：imgsz=640（图像resize尺寸），epochs=50，batch=16，device=0（GPU））
results = model.train(
    data=os.path.join(yolo_data_dir, "data.yaml"),
    imgsz=640,
    epochs=50,
    batch=16,
    device=0,
    name="nucls_yolov8n_detection"
)

# 5. 验证模型性能（输出mAP50、召回率等指标）
val_results = model.val(data=os.path.join(yolo_data_dir, "data.yaml"), imgsz=640)
print(f"模型mAP50（目标检测关键指标）：{val_results.box.map50:.4f}")
print(f"肿瘤细胞核（tumor）召回率：{val_results.box.recall[2]:.4f}")  # recall[2]对应tumor类别（ID=2）