Utils/DatasetDownloader.py - master - rodion_1521/EntropyMapSegmentation

Utils/DatasetDownloader.py
155 строк6 KB
rodion_1521
benchmark
02 июн 2025, 11:32
02 июн 2025, 11:329a12897
import os
import shutil
from pathlib import Path
import numpy as np
from PIL import Image
import scipy.stats
from torchvision.datasets import VOCSegmentation
from pycocotools.coco import COCO
import urllib.request
import tarfile
import zipfile
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def compute_background_entropy(image, mask, bg_label=0):
    try:
        image_np = np.array(image.convert('L'))
        mask_np = np.array(mask)
        bg_mask = (mask_np == bg_label)
        bg_pixels = image_np[bg_mask]
        if len(bg_pixels) == 0:
            return float('inf')
        hist, _ = np.histogram(bg_pixels, bins=256, range=(0, 255), density=True)
        entropy = scipy.stats.entropy(hist + 1e-10)
        return entropy
    except Exception as e:
        logging.error(f"Ошибка при вычислении энтропии: {e}")
        return float('inf')


def download_voc2012(data_dir):
    """Скачивает и распаковывает PASCAL VOC 2012."""
    voc_dir = Path(data_dir) / 'VOC2012'
    if voc_dir.exists():
        logging.info("PASCAL VOC 2012 уже загружен.")
        return voc_dir
    try:
        dataset = VOCSegmentation(root=data_dir, year='2012', image_set='trainval', download=True)
        logging.info("PASCAL VOC 2012 успешно загружен.")
        return voc_dir
    except Exception as e:
        logging.error(f"Ошибка при загрузке VOC2012: {e}")
        return None


def download_coco2017(data_dir):
    """Скачивает и распаковывает MS COCO 2017 (train2017 и аннотации)."""
    coco_dir = Path(data_dir) / 'COCO2017'
    train_dir = coco_dir / 'train2017'
    ann_dir = coco_dir / 'annotations'

    if train_dir.exists() and ann_dir.exists() and (ann_dir / 'instances_train2017.json').exists():
        logging.info("MS COCO 2017 уже загружен (найдены локальные файлы).")
        return coco_dir

    try:
        coco_dir.mkdir(parents=True, exist_ok=True)
        img_url = 'http://images.cocodataset.org/zips/train2017.zip'
        ann_url = 'http://images.cocodataset.org/annotations/annotations_trainval2017.zip'
        img_zip = coco_dir / 'train2017.zip'
        ann_zip = coco_dir / 'annotations.zip'

        # Скачивание изображений
        logging.info("Скачивание изображений COCO train2017...")
        try:
            urllib.request.urlretrieve(img_url, img_zip)
        except urllib.error.HTTPError as e:
            return None
        with zipfile.ZipFile(img_zip, 'r') as zip_ref:
            zip_ref.extractall(coco_dir)
        img_zip.unlink()

        logging.info("Скачивание аннотаций COCO...")
        try:
            urllib.request.urlretrieve(ann_url, ann_zip)
        except urllib.error.HTTPError as e:
            return None
        with zipfile.ZipFile(ann_zip, 'r') as zip_ref:
            zip_ref.extractall(coco_dir)
        ann_zip.unlink()

        logging.info("MS COCO 2017 успешно загружен.")
        return coco_dir
    except Exception as e:
        logging.error(f"Ошибка при загрузке COCO2017: {e}")
        return None


def filter_dataset(data_dir, entropy_threshold=2.0, max_images=50):
    output_dir = Path(data_dir) / 'dataset'
    output_dir.mkdir(parents=True, exist_ok=True)
    filtered_count = 0
    entropies = []

    voc_dir = download_voc2012(data_dir)
    if voc_dir:
        logging.info("Фильтрация PASCAL VOC 2012...")
        dataset = VOCSegmentation(root=data_dir, year='2012', image_set='trainval', download=False)
        for idx in range(len(dataset)):
            if filtered_count >= max_images:
                break
            image, mask = dataset[idx]
            entropy = compute_background_entropy(image, mask)
            if entropy < entropy_threshold:
                image.save(output_dir / f'voc_image_{idx}.jpg')
                mask.save(output_dir / f'voc_mask_{idx}.png')
                filtered_count += 1
                entropies.append(entropy)
                logging.info(f"VOC: Сохранено изображение {idx}, энтропия: {entropy:.2f}")

    coco_dir = download_coco2017(data_dir)
    if coco_dir:
        logging.info("Фильтрация MS COCO 2017...")
        coco = COCO(coco_dir / 'annotations' / 'instances_train2017.json')
        img_ids = coco.getImgIds()
        for img_id in img_ids:
            if filtered_count >= max_images:
                break
            img_info = coco.loadImgs(img_id)[0]
            img_path = coco_dir / 'train2017' / img_info['file_name']
            try:
                image = Image.open(img_path).convert('RGB')
            except FileNotFoundError:
                logging.warning(f"Изображение {img_path} не найдено, пропускаем.")
                continue
            ann_ids = coco.getAnnIds(imgIds=img_id, iscrowd=False)
            if not ann_ids:
                continue
            anns = coco.loadAnns(ann_ids)
            mask = np.zeros((img_info['height'], img_info['width']), dtype=np.uint8)
            for ann in anns:
                mask |= coco.annToMask(ann)  # Объединяем маски всех объектов
            mask = Image.fromarray(mask * 255)  # Преобразуем в изображение
            entropy = compute_background_entropy(image, mask, bg_label=0)
            if entropy < entropy_threshold:
                image.save(output_dir / f'coco_image_{img_id}.jpg')
                mask.save(output_dir / f'coco_mask_{img_id}.png')
                filtered_count += 1
                entropies.append(entropy)
                logging.info(f"COCO: Сохранено изображение {img_id}, энтропия: {entropy:.2f}")
    else:
        logging.warning("MS COCO 2017 не загружен, продолжаем только с VOC.")

    logging.info(f"Отфильтровано {filtered_count} изображений.")
    if entropies:
        logging.info(f"Средняя энтропия фона: {np.mean(entropies):.2f}, Мин: {np.min(entropies):.2f}, Макс: {np.max(entropies):.2f}")


if __name__ == "__main__":
    data_dir = './datasets'
    entropy_threshold = 2.0  # Порог энтропии для однородного фона
    max_images = 50  # Максимальное количество изображений для сохранения
    filter_dataset(data_dir, entropy_threshold, max_images)