from tqdm import tqdm
import detectron2.data.transforms as T
from detectron2.config import LazyCall as L
from detectron2.data import (
    build_detection_test_loader,
    build_detection_train_loader,
    get_detection_dataset_dicts,
    filter_images_with_only_crowd_annotations
)
from detectron2.evaluation import COCOEvaluator
from detrex.data import DetrDatasetMapper
from omegaconf import OmegaConf
import os
import json
from PIL import Image
import logging

def crop_and_save_instances_with_filter(dataset_dicts, output_dir="cropped_instances"):
    """
    从数据集字典中过滤无效标注和裁剪每个实例并保存。

    Args:
        dataset_dicts: 数据集字典列表。
        output_dir (str): 保存裁剪图像和标注的目录。
    """
    # 使用 filter_images_with_only_crowd_annotations 过滤数据
    dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)

    os.makedirs(output_dir, exist_ok=True)
    cropped_annotations = []

    # 使用 tqdm 显示进度条
    with tqdm(total=len(dataset_dicts), desc="Cropping Instances", unit="image") as pbar:
        for idx, sample in enumerate(dataset_dicts):
            file_name = sample["file_name"]
            try:
                image = Image.open(file_name).convert("RGB")
            except Exception as e:
                logging.warning(f"Skipping corrupted image {file_name}: {e}")
                pbar.update(1)
                continue

            annotations = sample["annotations"]

            for ann_idx, annotation in enumerate(annotations):
                bbox = annotation["bbox"]  # [x, y, w, h]
                category_id = annotation["category_id"]

            # 检查边界框合法性
            x1, y1, w, h = bbox
            if w <= 0 or h <= 0:
                logging.warning(f"Skipping invalid bbox with zero area: {bbox}")
                continue

            x2, y2 = x1 + w, y1 + h

            # 裁剪图像
            cropped_img = image.crop((int(x1), int(y1), int(x2), int(y2)))

            # 确保裁剪后的图像非空
            if cropped_img.size[0] <= 0 or cropped_img.size[1] <= 0:
                logging.warning(f"Skipping empty cropped image for bbox: {bbox}")
                continue

            cropped_file_name = f"{idx}_{ann_idx}_class_{category_id}.jpg"
            cropped_img_path = os.path.join(output_dir, cropped_file_name)
            cropped_img.save(cropped_img_path)

            # 保存标注信息
            cropped_annotations.append({
                "file_name": cropped_file_name,
                "width": int(x2 - x1),
                "height": int(y2 - y1),
                "category_id": category_id,
                "original_file": file_name,
                "original_bbox": bbox
            })

            # 更新进度条
            pbar.update(1)

    # 保存标注文件
    annotation_file = os.path.join(output_dir, "annotations.json")
    with open(annotation_file, "w") as f:
        json.dump(cropped_annotations, f, indent=4)
    logging.info(f"Annotations saved to {annotation_file}")
    logging.info(f"All cropped instances saved to {output_dir}")


# 加载训练数据集并裁剪
train_dataset_dicts = get_detection_dataset_dicts(names="coco_2017_train_ovd_unipro")
crop_and_save_instances_with_filter(train_dataset_dicts, output_dir="cropped_train_instances")

# 加载测试数据集并裁剪
test_dataset_dicts = get_detection_dataset_dicts(names="coco_2017_val_ovd", filter_empty=False)
crop_and_save_instances_with_filter(test_dataset_dicts, output_dir="cropped_test_instances")

# 替换 train_dataset_dicts 和 test_dataset_dicts
cropped_train_annotations = "cropped_train_instances/annotations.json"
cropped_test_annotations = "cropped_test_instances/annotations.json"

with open(cropped_train_annotations, "r") as f:
    train_dataset_dicts = json.load(f)

with open(cropped_test_annotations, "r") as f:
    test_dataset_dicts = json.load(f)


# 构建数据加载器配置
dataloader = OmegaConf.create()

dataloader.train = L(build_detection_train_loader)(
    dataset=train_dataset_dicts,  # 使用裁剪后的训练数据集
    mapper=L(DetrDatasetMapper)(
        augmentation=[
            L(T.RandomFlip)(),
            L(T.ResizeShortestEdge)(
                short_edge_length=(
                    480,
                    512,
                    544,
                    576,
                    608,
                    640,
                    672,
                    704,
                    736,
                    768,
                    800,
                ),
                max_size=1333,
                sample_style="choice",
            ),
        ],
        augmentation_with_crop=[
            L(T.RandomFlip)(),
            L(T.ResizeShortestEdge)(
                short_edge_length=(400, 500, 600),
                sample_style="choice",
            ),
            L(T.RandomCrop)(
                crop_type="absolute_range",
                crop_size=(384, 600),
            ),
            L(T.ResizeShortestEdge)(
                short_edge_length=(
                    480,
                    512,
                    544,
                    576,
                    608,
                    640,
                    672,
                    704,
                    736,
                    768,
                    800,
                ),
                max_size=1333,
                sample_style="choice",
            ),
        ],
        is_train=True,
        mask_on=False,
        img_format="RGB",
    ),
    total_batch_size=16,
    num_workers=2,
)

dataloader.test = L(build_detection_test_loader)(
    dataset=test_dataset_dicts,  # 使用裁剪后的测试数据集
    mapper=L(DetrDatasetMapper)(
        augmentation=[
            L(T.ResizeShortestEdge)(
                short_edge_length=800,
                max_size=1333,
            ),
        ],
        augmentation_with_crop=None,
        is_train=False,
        mask_on=False,
        img_format="RGB",
    ),
    num_workers=2,
)

dataloader.evaluator = L(COCOEvaluator)(
    # dataset_name="${..test.dataset.names}",
    dataset_name="coco_2017_val_ovd",
)
