Files
ai-med-edu/seminars/FeatureExtractionClient.ipynb
T

22 KiB
Raw Blame History

In [ ]:
!pip install pandas
In [ ]:
!pip install pillow
In [ ]:
!pip install torch==2.6.0 torchaudio==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
In [ ]:
In [ ]:
#!/usr/bin/env python3
# =============================================================================
# FeatureService Client — Hybrid Version
# Батчинг + поллинг статуса + пост-обработка CSV
# =============================================================================

import os
import csv
import time
import requests
import pandas as pd
from PIL import Image
import numpy as np
import io
import re
from http import HTTPStatus
from pathlib import Path

# =============================================================================
# CONFIG
# =============================================================================

# Сервер
#SERVER_IP = "85.143.113.40"
SERVER_IP = "192.168.203.91"
#SERVER_IP = "127.0.0.1"
SERVER_PORT = 14407
BASE_URL = f"http://{SERVER_IP}:{SERVER_PORT}"

# Пути
DIR_PATH = "../data/"
RESEARCH_DIR = "prep/"
PATH_IMAGES = DIR_PATH + RESEARCH_DIR + "images"
PATH_MASKS = DIR_PATH + RESEARCH_DIR + "images_results/masks_sam"
PATH_CSV = DIR_PATH + RESEARCH_DIR + "csv"
PATH_CLASSES = DIR_PATH +  RESEARCH_DIR + "predefined_classes.txt"

# Форматы
IMAGE_EXT = ".jpg"
MASK_EXT = ".png"

# Параметры обработки
BATCH_SIZE = 2              # Сколько задач отправлять параллельно
POLL_INTERVAL = 1           # Интервал опроса статуса (сек)
POLL_TIMEOUT = 300          # Максимальное время ожидания задачи (сек)
REQUEST_TIMEOUT = 120       # Таймаут HTTP-запроса (сек)
DELAY_BETWEEN_TASKS = 0.3   # Пауза между задачами внутри батча

# =============================================================================
# GLOBAL STATE
# =============================================================================

in_progress = []  # Очередь задач: [(image, mask, task_id), ...]
processed = []    # Успешно обработанные: [(image, mask, task_id, output_csv), ...]

# =============================================================================
# UTILS
# =============================================================================

def ensure_dirs():
    """Создаёт необходимые директории"""
    for d in [PATH_IMAGES, PATH_MASKS, PATH_CSV]:
        os.makedirs(d, exist_ok=True)

def load_classes():
    """Загружает словарь классов из файла"""
    classes = {}
    try:
        with open(PATH_CLASSES, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                parts = line.split(maxsplit=1)
                if len(parts) == 2:
                    classes[int(parts[0])] = parts[1].strip()
    except FileNotFoundError:
        print(f"⚠️ {PATH_CLASSES} не найден, метки будут 'Unknown'")
    except Exception as e:
        print(f"⚠️ Ошибка загрузки классов: {e}")
    return classes

def get_label(filename, classes):
    """Извлекает метку из имени файла: _class_7_ → класс 7 → название"""
    match = re.search(r'_class_(\d+)_', filename)
    if match:
        class_id = int(match.group(1))
        return classes.get(class_id, f"Class_{class_id}")
    return "Unknown"

def normalize_mask(img):
    """Нормализует маску: все ненулевые пиксели → 1 (для PyRadiomics)"""
    arr = np.array(img)
    arr[arr > 0] = 1
    return Image.fromarray(arr.astype(np.uint8), mode='L')

def match_mask_to_image(image_short_name, masks_list):
    """
    Точное сопоставление маски с изображением.
    Ищет маски, начинающиеся с имени изображения.
    """
    matched = []
    for mask in masks_list:
        if mask.startswith(image_short_name) and mask.endswith(MASK_EXT):
            matched.append(mask)
    return matched

# =============================================================================
# CORE: MODEL COMMUNICATION
# =============================================================================

def launch_model(short_name, image_name, mask_name):
    """
    Отправляет изображение и маску на сервер.
    Возвращает task_id или None при ошибке.
    """
    image_path = os.path.join(PATH_IMAGES, image_name)
    mask_path = os.path.join(PATH_MASKS, mask_name)
    
    if not os.path.exists(image_path):
        print(f"✗ Файл не найден: {image_path}")
        return None
    if not os.path.exists(mask_path):
        print(f"✗ Файл не найден: {mask_path}")
        return None

    print(f"🔄 Отправка: {image_name} + {mask_name}")
    
    try:
        # Читаем изображение
        with open(image_path, 'rb') as img_f:
            image_bytes = img_f.read()
        
        # Нормализуем и читаем маску
        with Image.open(mask_path) as mask_img:
            normalized = normalize_mask(mask_img)
            buf = io.BytesIO()
            normalized.save(buf, format='PNG')
            buf.seek(0)
            mask_bytes = buf.read()
        
        # Формируем multipart/form-data с правильными MIME-типами
        files = [
            ('image', ('image.jpg', image_bytes, 'image/jpeg')),
            ('mask1', ('mask.png', mask_bytes, 'image/png')),
        ]
        
        resp = requests.post(
            f'{BASE_URL}/texture/feature', 
            files=files, 
            timeout=REQUEST_TIMEOUT
        )
        resp.raise_for_status()
        
        task_id = resp.text.strip().strip('"').strip("'")
        if not task_id or task_id.lower() == 'error':
            print(f"✗ Сервер вернул ошибку: {resp.text}")
            return None
            
        print(f"✓ Задача создана: {task_id}")
        return task_id
        
    except requests.exceptions.RequestException as e:
        print(f"✗ Сетевая ошибка: {e}")
        return None
    except Exception as e:
        print(f"✗ Ошибка обработки: {e}")
        return False

def wait_for_result(task_id, mask_name):
    """
    Поллинг статуса задачи до завершения.
    Сохраняет результат в CSV при успехе.
    Возвращает путь к сохранённому файлу или None.
    """
    mask_base = os.path.splitext(mask_name)[0]
    output_csv = os.path.join(PATH_CSV, f"{mask_base}.csv")
    
    start_time = time.time()
    
    while True:
        # Проверка таймаута
        if time.time() - start_time > POLL_TIMEOUT:
            print(f"✗ Таймаут ожидания для {task_id}")
            return None
        
        try:
            status_resp = requests.get(
                f'{BASE_URL}/texture/status?id={task_id}',
                timeout=REQUEST_TIMEOUT
            )
            
            if status_resp.status_code == HTTPStatus.OK:
                print(f"✓ Задача {task_id} завершена, получаем результат...")
                
                result_resp = requests.get(
                    f'{BASE_URL}/texture/result?id={task_id}',
                    timeout=REQUEST_TIMEOUT
                )
                
                if result_resp.status_code == 200 and 'id not found' not in result_resp.text.lower():
                    os.makedirs(PATH_CSV, exist_ok=True)
                    with open(output_csv, "w", encoding="utf-8") as f:
                        f.write(result_resp.text)
                    print(f"💾 Сохранено: {output_csv}")
                    return output_csv
                else:
                    print(f"✗ Ошибка получения результата: {result_resp.text[:100]}")
                    return None
            else:
                # Задача ещё в процессе
                time.sleep(POLL_INTERVAL)
                
        except requests.exceptions.RequestException as e:
            print(f"⚠️ Ошибка опроса статуса: {e}")
            time.sleep(POLL_INTERVAL)
        except Exception as e:
            print(f"✗ Неожиданная ошибка: {e}")
            return None

def process_batch():
    """
    Обрабатывает очередь in_progress: ждёт завершения всех задач и сохраняет результаты.
    """
    global in_progress, processed
    
    if not in_progress:
        return
    
    print(f"\n⏳ Ожидание завершения {len(in_progress)} задач...")
    
    for image_name, mask_name, task_id in in_progress:
        output_path = wait_for_result(task_id, mask_name)
        if output_path:
            processed.append((image_name, mask_name, task_id, output_path))
    
    in_progress.clear()
    print("✅ Батч обработан\n")

# =============================================================================
# POST-PROCESSING: CSV MERGE & ENRICH
# =============================================================================

def delete_useless_columns(row):
    """Удаляет служебные колонки PyRadiomics"""
    diagnostics_cols = [
        'diagnostics_Versions_PyRadiomics', 'diagnostics_Versions_Numpy',
        'diagnostics_Versions_SimpleITK', 'diagnostics_Versions_PyWavelet',
        'diagnostics_Versions_Python', 'diagnostics_Configuration_Settings',
        'diagnostics_Configuration_EnabledImageTypes', 'diagnostics_Image-original_Hash',
        'diagnostics_Image-original_Dimensionality', 'diagnostics_Image-original_Spacing',
        'diagnostics_Image-original_Size', 'diagnostics_Mask-original_CenterOfMassIndex',
        'diagnostics_Mask-original_CenterOfMass', 'diagnostics_Mask-original_BoundingBox',
        'diagnostics_Mask-original_Size', 'diagnostics_Mask-original_Hash',
        'diagnostics_Mask-original_Spacing'
    ]
    for col in diagnostics_cols:
        row.pop(col, None)
    return row

def merge_csv_files():
    """Объединяет все CSV в result.csv"""
    csv_files = [f for f in os.listdir(PATH_CSV) 
                 if f.endswith('.csv') and f != 'result.csv']
    if not csv_files:
        print("⚠️ Нет CSV-файлов для объединения")
        return None
    
    print(f"🔗 Объединяю {len(csv_files)} файлов...")
    all_rows = []
    
    for csv_file in csv_files:
        try:
            with open(os.path.join(PATH_CSV, csv_file), 'r', newline='', encoding='utf-8') as f:
                for row in csv.DictReader(f):
                    row = delete_useless_columns(row)
                    row['source_file'] = csv_file
                    all_rows.append(row)
        except Exception as e:
            print(f"⚠️ Ошибка чтения {csv_file}: {e}")
    
    if not all_rows:
        print("⚠️ Нет данных для сохранения")
        return None
    
    result_path = os.path.join(PATH_CSV, 'result.csv')
    pd.DataFrame(all_rows).to_csv(result_path, index=False, encoding='utf-8')
    print(f"💾 Объединённый файл: {result_path}")
    return result_path

def add_label_and_path_to_cell(result_path=None):
    """
    Добавляет label и path_to_cell ПЕРВЫМИ двумя столбцами в result.csv
    """
    if result_path is None:
        result_path = os.path.join(PATH_CSV, 'result.csv')
    
    if not os.path.exists(result_path):
        print(f"⚠️ Файл не найден: {result_path}")
        return
    
    classes = load_classes()
    rows = []
    counter = 1
    
    print("🏷️  Добавляю label и path_to_cell...")
    
    with open(result_path, 'r', newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        
        for row in reader:
            # Ищем имя файла в разных возможных полях
            filename = (row.get('name') or row.get('source_file') or 
                       row.get('path_to_file') or row.get('filename') or '')
            
            # Извлекаем label
            if filename:
                basename = os.path.splitext(os.path.basename(filename))[0]
                label = get_label(basename, classes)
            else:
                label = "Unknown"
            
            # Формируем path_to_cell: "1_F0000001_class_7_obj_0000.png.csv"
            source = row.get('source_file', os.path.basename(filename))
            path_to_cell = f"{counter}_{source}"
            
            # Создаём новый dict с label и path_to_cell первыми
            new_row = {'label': label, 'path_to_cell': path_to_cell}
            for key, value in row.items():
                if key not in ('label', 'path_to_cell'):
                    new_row[key] = value
            
            rows.append(new_row)
            counter += 1
    
    if rows:
        final_columns = ['label', 'path_to_cell'] + [c for c in rows[0].keys() if c not in ['label', 'path_to_cell']]
        pd.DataFrame(rows).to_csv(
            result_path, 
            index=False, 
            encoding='utf-8',
            columns=final_columns
        )
        print(f"✅ Готово: {result_path}")
        return result_path
    else:
        print("⚠️ Нет строк для обработки")
        return None

# =============================================================================
# MAIN PROCESSING LOOP
# =============================================================================

def process_all_images():
    """
    Основной цикл: сканирует изображения, сопоставляет маски,
    отправляет задачи батчами с поллингом статуса.
    """
    global in_progress, counter
    
    ensure_dirs()
    
    # Сканируем файлы
    images = [f for f in os.listdir(PATH_IMAGES) if f.lower().endswith(IMAGE_EXT)]
    masks = [f for f in os.listdir(PATH_MASKS) if f.lower().endswith(MASK_EXT)]
    
    if not images:
        print(f"⚠️ Нет изображений в {PATH_IMAGES}")
        return
    if not masks:
        print(f"⚠️ Нет масок в {PATH_MASKS}")
        return
    
    print(f"📁 Найдено: {len(images)} изображений, {len(masks)} масок")
    
    counter = 0  # Счётчик задач в текущем батче
    
    for img in images:
        short = os.path.splitext(img)[0]
        img_masks = match_mask_to_image(short, masks)
        
        if not img_masks:
            print(f"⚠️ Нет масок для {img}")
            continue
        
        for mask in img_masks:
            # Если батч заполнен — обрабатываем его
            if counter >= BATCH_SIZE:
                process_batch()
                counter = 0
            
            # Отправляем задачу
            task_id = launch_model(short, img, mask)
            if task_id:
                in_progress.append((img, mask, task_id))
                counter += 1
                time.sleep(DELAY_BETWEEN_TASKS)  # Не перегружаем сервер
    
    # Обрабатываем оставшиеся задачи
    if in_progress:
        process_batch()
    
    print(f"\n🎯 Всего обработано: {len(processed)} задач")

# =============================================================================
# ENTRY POINT
# =============================================================================

if __name__ == "__main__":
    print(f"🚀 FeatureService Client (Hybrid)")
    print(f"🌐 Сервер: {BASE_URL}")
    print(f"📦 BATCH_SIZE: {BATCH_SIZE}, ⏱️ POLL_TIMEOUT: {POLL_TIMEOUT}s\n")
    
    try:
        # 1. Обработка изображений
        print("=== [1/3] Отправка задач на сервер ===")
        process_all_images()
        
        if not processed:
            print("⚠️ Нет успешно обработанных задач. Пропускаю пост-обработку.")
        else:
            # 2. Объединение CSV
            print("\n=== [2/3] Объединение результатов ===")
            result_path = merge_csv_files()
            
            # 3. Добавление меток
            if result_path:
                print("\n=== [3/3] Добавление label и path_to_cell ===")
                add_label_and_path_to_cell(result_path)
        
        print(f"\n✨ Готово. Результат: {PATH_CSV}/result.csv")
        
    except KeyboardInterrupt:
        print("\n⚠️ Прервано пользователем")
    except Exception as e:
        print(f"\n✗ Критическая ошибка: {e}")
        import traceback
        traceback.print_exc()
In [ ]: