new_file - master - desskas/chapter2

new_file
89 строк3 KB
desskas
update: new_file
26 май 2026, 17:44
26 май 2026, 17:44c75c7d8
"""
Скрипт подготовки датасета для курсовой работы (Глава 2)
Задача: Бинарная классификация Cat vs Dog
Студент: Уразова Д.В., ЕТ-115
"""

import os
import shutil
import random
from PIL import Image

#НАСТРОЙКИ
RANDOM_SEED = 42
SAMPLE_PER_CLASS = 300
TRAIN_RATIO = 0.8
MIN_SIZE = 512
TARGET_SIZE = (512, 512)

BASE_PATH = 'data'
SOURCE_CAT = os.path.join(BASE_PATH, 'raw', 'Cat')
SOURCE_DOG = os.path.join(BASE_PATH, 'raw', 'Dog')

#1. СОЗДАНИЕ ПОДВЫБОРКИ
def prepare_split():
    random.seed(RANDOM_SEED)
    
    def get_images(folder):
        return [f for f in os.listdir(folder) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    
    def copy_files(src, dst_folder, files):
        os.makedirs(dst_folder, exist_ok=True)
        for f in files:
            shutil.copy2(os.path.join(src, f), os.path.join(dst_folder, f))

    # Кошки
    cats = get_images(SOURCE_CAT)
    random.shuffle(cats)
    selected_cats = cats[:SAMPLE_PER_CLASS]
    split_c = int(len(selected_cats) * TRAIN_RATIO)
    copy_files(SOURCE_CAT, 'data/train/cat', selected_cats[:split_c])
    copy_files(SOURCE_CAT, 'data/val/cat', selected_cats[split_c:])
    
    # Собаки
    dogs = get_images(SOURCE_DOG)
    random.shuffle(dogs)
    selected_dogs = dogs[:SAMPLE_PER_CLASS]
    split_d = int(len(selected_dogs) * TRAIN_RATIO)
    copy_files(SOURCE_DOG, 'data/train/dog', selected_dogs[:split_d])
    copy_files(SOURCE_DOG, 'data/val/dog', selected_dogs[split_d:])
    
    print(" Подвыборка создана: 600 изображений")

#2. ПРОВЕРКА КАЧЕСТВА 
def check_quality():
    folders = ['data/train/cat', 'data/train/dog', 'data/val/cat', 'data/val/dog']
    print("\n Проверка качества:")
    for path in folders:
        total, broken, small = 0, 0, 0
        for f in os.listdir(path):
            if not f.lower().endswith(('.jpg','.jpeg','.png')): continue
            total += 1
            try:
                w, h = Image.open(os.path.join(path, f)).size
                if w < MIN_SIZE or h < MIN_SIZE: small += 1
            except: broken += 1
        print(f"  {path}: всего={total}, битых={broken}, <512px={small}")

#3. МАСШТАБИРОВАНИЕ
def resize_all():
    print("\nМасштабирование до 512x512...")
    folders = ['data/train/cat', 'data/train/dog', 'data/val/cat', 'data/val/dog']
    count = 0
    for path in folders:
        for f in os.listdir(path):
            if not f.lower().endswith(('.jpg','.jpeg','.png')): continue
            fp = os.path.join(path, f)
            img = Image.open(fp)
            if img.width < MIN_SIZE or img.height < MIN_SIZE:
                img = img.resize(TARGET_SIZE, Image.LANCZOS)
                img.save(fp)
                count += 1
    print(f"Увеличено изображений: {count}")

# ЗАПУСК
if __name__ == "__main__":
    prepare_split()
    check_quality()      # Проверка до увеличения
    resize_all()         # Увеличение
    check_quality()      # Проверка после увеличения