diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e69de29 diff --git a/inference.py b/inference.py index 1a31822..2e0b3f0 100644 --- a/inference.py +++ b/inference.py @@ -5,7 +5,10 @@ import matplotlib.pyplot as plt import argparse from src import UNet -from src.quality.quality_scorer_old import QualityScorer +from src.quality.quality_scorer import QualityScorer + + +# from src.quality.quality_scorer_old import QualityScorer def load_model(model_path, in_channels=3, out_classes=2, device='cpu'): diff --git a/main.py b/main.py index 8f8ca2f..ef90058 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,4 @@ -# Press the green button in the gutter to run the script. from pydicom_custom import main if __name__ == '__main__': main() -# See PyCharm help at https://www.jetbrains.com/help/pycharm/ diff --git a/src/__init__.py b/src/__init__.py index 9954658..841f5d2 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -11,8 +11,6 @@ __author__ = "Your Team" from src.config import Config from src.data_loader import FlexibleDataset from src.model import SegmentationEngine, UNet -from src.quality import QualityScorer -from src.api import app # Что импортируется при "from src import *" __all__ = [ @@ -21,11 +19,12 @@ __all__ = [ "UNet", "SegmentationEngine", "QualityScorer", - "app", "__version__", ] # Можно добавить инициализацию логгера import logging +from src.quality.quality_scorer import QualityScorer + logging.getLogger(__name__).addHandler(logging.NullHandler()) \ No newline at end of file diff --git a/src/api/__init__.py b/src/api/__init__.py index 7aa7f09..53371f2 100644 --- a/src/api/__init__.py +++ b/src/api/__init__.py @@ -2,9 +2,10 @@ REST API для сервиса оценки качества """ -from src.api.endpoints import app +from .endpoints import router as endpoints +from .root import router as root -# Можно добавить middleware или настройки __all__ = [ - "app", + "endpoints", + "root" ] \ No newline at end of file diff --git a/src/api/endpoints.py b/src/api/endpoints.py index b847d24..644300d 100644 --- a/src/api/endpoints.py +++ b/src/api/endpoints.py @@ -1,38 +1,23 @@ -from fastapi import FastAPI, File, UploadFile, APIRouter -from fastapi.responses import JSONResponse, FileResponse -from fastapi.staticfiles import StaticFiles -from pydantic import BaseModel +from fastapi import File, UploadFile, APIRouter import numpy as np from PIL import Image import io import torch from pathlib import Path -from typing import List, Optional, Dict, Any import base64 from src import UNet +from src.api.schemas.health_response import HealthResponse +from src.api.schemas.quality_response import QualityResponse +from src.core import orchestrator +from src.core.config import config +from src.core.orchestrator import Orchestrator from src.quality.quality_scorer import QualityScorer, convert_to_serializable +from src.utils.utils import get_device -# Создаем приложение -app = FastAPI( - title="Bone Quality Assessment API", - description="API для оценки качества медицинских изображений", - version="1.0.0", - docs_url="/docs", - redoc_url="/redoc", - openapi_url="/openapi.json" -) +orchestrator = Orchestrator(config) -# Подключаем статические файлы -static_dir = Path(__file__).parent / "static" -if static_dir.exists(): - app.mount("/static", StaticFiles(directory=str(static_dir)), name="static") - print(f"✅ Статика подключена: {static_dir}") -else: - print(f"⚠️ Папка static не найдена: {static_dir}") - -# Создаем роутер -router = APIRouter(prefix="/api/v1", tags=["quality"]) +router = APIRouter(prefix="/api/v1") # Глобальные переменные model = None @@ -40,16 +25,6 @@ scorer = None device = None -def get_device(): - """Определение устройства""" - if torch.backends.mps.is_available(): - return torch.device("mps") - elif torch.cuda.is_available(): - return torch.device("cuda") - else: - return torch.device("cpu") - - def load_models(): """Ленивая загрузка моделей""" global model, scorer, device @@ -81,25 +56,6 @@ def load_models(): return model, scorer, device -class QualityResponse(BaseModel): - overall_quality: str - severity: str - issues: List[Dict[str, Any]] - confidence: float - metrics: Dict[str, Any] - position_validation: Optional[Dict[str, Any]] = None - artifact_validation: Optional[Dict[str, Any]] = None - mask: Optional[str] = None # base64 encoded mask image - - -class HealthResponse(BaseModel): - status: str - model_loaded: bool - device: str - scorer_loaded: bool - model_path_exists: bool - - @router.post("/analyze", response_model=QualityResponse) async def analyze_image(file: UploadFile = File(...)): """Анализ изображения на качество""" @@ -159,6 +115,8 @@ async def analyze_image(file: UploadFile = File(...)): # Оценка качества quality_scores = scorer.evaluate(image=original_image, segmentation=mask) + result = orchestrator.analyze(image) + print(f"\n📊 Результаты:") print(f" Качество: {quality_scores['overall_quality']}") print(f" Уверенность: {quality_scores['confidence']:.2f}") @@ -210,58 +168,13 @@ async def health_check(): ) -# Добавляем роутер в приложение -app.include_router(router) +@router.post("/switch_mode") +async def switch_mode(mode: str): + """Переключение режима работы (разработка/медицина)""" + orchestrator.switch_mode(mode) + return {"status": "ok", "mode": mode} - -# Корневой эндпоинт -@app.get("/") -async def root(): - """Главная страница""" - index_path = static_dir / "index.html" - if index_path.exists(): - return FileResponse(str(index_path)) - else: - return { - "service": "Bone Quality Assessment API", - "version": "1.0.0", - "docs": "/docs", - "health": "/api/v1/health", - "analyze": "/api/v1/analyze (POST)" - } - - -# Обработчик 404 -@app.exception_handler(404) -async def not_found_handler(request, exc): - return JSONResponse( - status_code=404, - content={ - "error": "Not Found", - "message": f"Endpoint {request.url.path} not found", - "available_endpoints": [ - "/", - "/docs", - "/redoc", - "/openapi.json", - "/api/v1/health", - "/api/v1/analyze (POST)", - "/static/" - ] - } - ) - - -if __name__ == "__main__": - import uvicorn - - print("🚀 Запуск Bone Quality Assessment API") - print(f"📍 Статика: {static_dir}") - print("🌐 http://localhost:8000") - uvicorn.run( - "src.api.endpoints:app", - host="0.0.0.0", - port=8000, - reload=True, - log_level="info" - ) \ No newline at end of file +@router.get("/info") +async def get_info(): + """Информация о текущем пайплайне""" + return orchestrator.get_info() diff --git a/src/api/root.py b/src/api/root.py new file mode 100644 index 0000000..00746e4 --- /dev/null +++ b/src/api/root.py @@ -0,0 +1,21 @@ +from fastapi import APIRouter +from pathlib import Path +from starlette.responses import FileResponse + +router = APIRouter(tags=["root"]) +@router.get("/") +async def root(): + """Главная страница""" + static_dir = Path(__file__).parent.parent / "api/static" + index_path = static_dir / "index.html" + + if index_path.exists(): + return FileResponse(str(index_path)) + else: + return { + "service": "Bone Quality Assessment API", + "version": "1.0.0", + "docs": "/docs", + "health": "/api/v1/health", + "analyze": "/api/v1/analyze (POST)" + } \ No newline at end of file diff --git a/src/api/schemas/health_response.py b/src/api/schemas/health_response.py new file mode 100644 index 0000000..98af0b4 --- /dev/null +++ b/src/api/schemas/health_response.py @@ -0,0 +1,9 @@ +from pydantic import BaseModel + + +class HealthResponse(BaseModel): + status: str + model_loaded: bool + device: str + scorer_loaded: bool + model_path_exists: bool \ No newline at end of file diff --git a/src/api/schemas/quality_response.py b/src/api/schemas/quality_response.py new file mode 100644 index 0000000..8b2250f --- /dev/null +++ b/src/api/schemas/quality_response.py @@ -0,0 +1,14 @@ +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel + + +class QualityResponse(BaseModel): + overall_quality: str + severity: str + issues: List[Dict[str, Any]] + confidence: float + metrics: Dict[str, Any] + position_validation: Optional[Dict[str, Any]] = None + artifact_validation: Optional[Dict[str, Any]] = None + mask: Optional[str] = None # base64 encoded mask image \ No newline at end of file diff --git a/src/classifiers/__init__.py b/src/classifiers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/classifiers/base.py b/src/classifiers/base.py new file mode 100644 index 0000000..728c837 --- /dev/null +++ b/src/classifiers/base.py @@ -0,0 +1,29 @@ +from abc import ABC, abstractmethod +import numpy as np +from PIL import Image +from typing import Union, Dict, Any + + +class BaseClassifier(ABC): + """Базовый класс для всех классификаторов""" + + @abstractmethod + def classify(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, Any]: + """ + Классификация объекта на изображении + + Args: + image: входное изображение + + Returns: + Словарь с результатами классификации + """ + pass + + @abstractmethod + def get_info(self) -> Dict[str, Any]: + """Информация о модели""" + return { + "name": self.__class__.__name__, + "type": "classification" + } \ No newline at end of file diff --git a/src/classifiers/bone_condition_classifier.py b/src/classifiers/bone_condition_classifier.py new file mode 100644 index 0000000..73067f1 --- /dev/null +++ b/src/classifiers/bone_condition_classifier.py @@ -0,0 +1,276 @@ +# src/models/bone_condition_classifier.py +from typing import Dict, Any, Union, Optional, List +import torch +import torch.nn as nn +import torchvision.models as models +from PIL import Image +import numpy as np +from pathlib import Path + +from src.classifiers.base import BaseClassifier + +# Маппинг состояний костей +BONE_CONDITIONS = { + 0: "Normal", + 1: "Osteopenia", + 2: "Osteoporosis", + 3: "Severe Osteoporosis", + 4: "Fracture", + 5: "Degenerative Changes", + 6: "Arthritis", + 7: "Tumor", + 8: "Infection", + 9: "Post-surgical Changes", + 10: "Congenital Anomaly" +} + +SEVERITY_MAP = { + "Normal": "None", + "Osteopenia": "Mild", + "Osteoporosis": "Moderate", + "Severe Osteoporosis": "Severe", + "Fracture": "Acute", + "Degenerative Changes": "Chronic", + "Arthritis": "Chronic", + "Tumor": "Severe", + "Infection": "Severe", + "Post-surgical Changes": "Mild", + "Congenital Anomaly": "Moderate" +} + + +class BoneConditionClassifier(BaseClassifier): + """ + Классификатор состояния костной ткани + Используется для медицинской диагностики на основе денситометрических изображений + """ + + def __init__( + self, + model_path: Optional[str] = None, + device: str = 'cpu', + num_classes: int = 11, # 10 состояний + 1 норма + input_size: int = 224 + ): + """ + Args: + model_path: путь к обученной модели (опционально) + device: устройство для инференса ('cpu', 'mps', 'cuda') + num_classes: количество классов состояний + input_size: размер входного изображения + """ + self.device = device + self.input_size = input_size + self.num_classes = num_classes + + # Используем предобученный ResNet50 как бэкбон + self.model = models.resnet50(pretrained=True) + + # Заменяем последний слой на количество состояний + num_features = self.model.fc.in_features + self.model.fc = nn.Sequential( + nn.Dropout(0.5), + nn.Linear(num_features, 256), + nn.ReLU(), + nn.Dropout(0.3), + nn.Linear(256, num_classes) + ) + + # Загружаем веса если есть + self.loaded = False + if model_path and Path(model_path).exists(): + try: + self.model.load_state_dict(torch.load(model_path, map_location=device)) + self.loaded = True + print(f"✅ Загружен классификатор костей из {model_path}") + except Exception as e: + print(f"⚠️ Ошибка загрузки модели из {model_path}: {e}") + print(" Используем неподготовленную модель (будет работать плохо)") + else: + print("⚠️ Классификатор костей не найден") + print(f" Ожидается: {model_path}") + print(" Обучите модель на медицинских данных перед использованием") + + self.model.to(device) + self.model.eval() + + # Нормализация для медицинских изображений (можно адаптировать под DICOM) + self.mean = np.array([0.485, 0.456, 0.406]) + self.std = np.array([0.229, 0.224, 0.225]) + + print(f"📊 Модель на устройстве: {device}") + print(f"📋 Количество классов: {num_classes}") + + def preprocess(self, image: Union[np.ndarray, Image.Image]) -> torch.Tensor: + """ + Предобработка изображения для модели + + Args: + image: входное изображение + + Returns: + torch.Tensor: подготовленный тензор + """ + # Конвертация в PIL Image + if isinstance(image, np.ndarray): + image = Image.fromarray(image) + + # Ресайз + image = image.resize((self.input_size, self.input_size)) + + # Конвертация в массив и нормализация + image_array = np.array(image, dtype=np.float32) / 255.0 + + # Если изображение grayscale, конвертируем в RGB + if len(image_array.shape) == 2: + image_array = np.stack([image_array] * 3, axis=2) + elif image_array.shape[2] == 1: + image_array = np.concatenate([image_array] * 3, axis=2) + + # Изменение порядка каналов HWC -> CHW + image_array = image_array.transpose(2, 0, 1) + + # Нормализация + for i in range(3): + image_array[i] = (image_array[i] - self.mean[i]) / self.std[i] + + return torch.FloatTensor(image_array).unsqueeze(0).to(self.device) + + def predict(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, Any]: + """ + Предсказание состояния костной ткани + + Args: + image: изображение в формате numpy array или PIL Image + + Returns: + Словарь с результатами: + - condition_id: ID состояния + - condition_name: название состояния + - severity: серьезность состояния + - confidence: уверенность (0-1) + - all_probs: вероятности всех классов + - loaded: загружена ли модель + """ + # Проверка загрузки модели + if not self.loaded: + return { + "condition_id": -1, + "condition_name": "Unknown", + "severity": "Unknown", + "confidence": 0.0, + "loaded": False, + "error": "Model not loaded" + } + + try: + # Предобработка + image_tensor = self.preprocess(image) + + # Инференс + with torch.no_grad(): + outputs = self.model(image_tensor) + probabilities = torch.softmax(outputs, dim=1) + confidence, predicted = torch.max(probabilities, 1) + + # Формируем результат + condition_id = predicted.item() + confidence_score = confidence.item() + all_probs = probabilities.cpu().numpy().tolist()[0] + + condition_name = BONE_CONDITIONS.get(condition_id, "Unknown") + severity = SEVERITY_MAP.get(condition_name, "Unknown") + + # Дополнительные метрики для медицинского контекста + is_abnormal = condition_id != 0 + risk_level = "Low" + if condition_id in [2, 3, 7, 8]: # Остеопороз, опухоль, инфекция + risk_level = "High" + elif condition_id in [1, 4, 5, 6]: # Остеопения, перелом, дегенерация + risk_level = "Medium" + + return { + "condition_id": condition_id, + "condition_name": condition_name, + "severity": severity, + "confidence": confidence_score, + "all_probs": all_probs, + "is_abnormal": is_abnormal, + "risk_level": risk_level, + "loaded": True, + "error": None + } + + except Exception as e: + return { + "condition_id": -1, + "condition_name": "Error", + "severity": "Unknown", + "confidence": 0.0, + "loaded": self.loaded, + "error": str(e) + } + + def classify(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, Any]: + """Алиас для predict (для совместимости с BaseClassifier)""" + return self.predict(image) + + def predict_batch(self, images: List[Union[np.ndarray, Image.Image]]) -> List[Dict[str, Any]]: + """ + Предсказание для батча изображений + + Args: + images: список изображений + + Returns: + Список результатов для каждого изображения + """ + return [self.predict(img) for img in images] + + def get_info(self) -> Dict[str, Any]: + """Информация о классификаторе""" + return { + "name": self.__class__.__name__, + "type": "classification", + "target": "bone_conditions", + "num_classes": self.num_classes, + "loaded": self.loaded, + "device": self.device, + "model_architecture": "ResNet50 + Custom Head", + "pretrained": True, + "input_size": self.input_size, + "conditions": list(BONE_CONDITIONS.values()), + "severity_map": SEVERITY_MAP + } + + def get_recommendations(self, condition_id: int) -> List[str]: + """ + Получение рекомендаций на основе состояния + + Args: + condition_id: ID состояния + + Returns: + Список рекомендаций + """ + recommendations = { + 0: ["✅ Состояние в норме", "Продолжайте мониторинг"], + 1: ["⚠️ Начальные изменения плотности", "Рекомендуется контроль через 6 месяцев", + "Увеличьте потребление кальция"], + 2: ["🔴 Умеренное снижение плотности", "Рекомендуется контроль через 3 месяца", "Консультация эндокринолога", + "Препараты кальция и витамин D"], + 3: ["🚨 Критическое снижение плотности", "Немедленная консультация специалиста", "Интенсивная терапия", + "Мониторинг переломов"], + 4: ["🦴 Обнаружен перелом", "Иммобилизация", "Консультация травматолога", "Контрольная рентгенография"], + 5: ["⚙️ Дегенеративные изменения", "Физиотерапия", "Противовоспалительная терапия", + "Контроль через 6 месяцев"], + 6: ["🔄 Артрит", "Противовоспалительная терапия", "Физиотерапия", "Консультация ревматолога"], + 7: ["🧬 Подозрение на опухоль", "Срочная консультация онколога", "МРТ/КТ исследование", "Биопсия"], + 8: ["🦠 Подозрение на инфекцию", "Антибактериальная терапия", "Консультация инфекциониста", + "Контрольный анализ"], + 9: ["🔧 Послеоперационные изменения", "Контроль через 3 месяца", "Физиотерапия", "Наблюдение хирурга"], + 10: ["🧬 Врожденная аномалия", "Консультация генетика", "Индивидуальный план лечения", "Мониторинг развития"] + } + return recommendations.get(condition_id, ["Обратитесь к специалисту"]) + + diff --git a/src/classifiers/pet_breed_classifier.py b/src/classifiers/pet_breed_classifier.py new file mode 100644 index 0000000..0c6611e --- /dev/null +++ b/src/classifiers/pet_breed_classifier.py @@ -0,0 +1,174 @@ +# src/models/pet_breed_classifier.py +from typing import Dict, Any, Union, Optional +import torch +import torch.nn as nn +import torchvision.models as models +from PIL import Image +import numpy as np +from pathlib import Path + +from src.classifiers.base import BaseClassifier + +# Маппинг ID породы в название +BREED_MAP = { + 1: "Abyssinian", 2: "Bengal", 3: "Birman", 4: "Bombay", + 5: "British Shorthair", 6: "Egyptian Mau", 7: "Maine Coon", + 8: "Persian", 9: "Ragdoll", 10: "Russian Blue", 11: "Siamese", + 12: "Sphynx", 13: "American Bulldog", 14: "American Pit Bull Terrier", + 15: "American Staffordshire Terrier", 16: "Australian Shepherd", + 17: "Beagle", 18: "Border Collie", 19: "Boxer", 20: "Chihuahua", + 21: "Cocker Spaniel", 22: "Dachshund", 23: "Doberman Pinscher", + 24: "English Cocker Spaniel", 25: "German Shepherd", 26: "Golden Retriever", + 27: "Great Dane", 28: "Jack Russell Terrier", 29: "Labrador Retriever", + 30: "Poodle", 31: "Rottweiler", 32: "Siberian Husky", + 33: "Staffordshire Bull Terrier", 34: "Yorkshire Terrier", + 35: "Dalmatian", 36: "Pug", 37: "Shih Tzu" +} + +SPECIES_MAP = { + 1: "Cat", + 2: "Dog" +} + + +class PetBreedClassifier(BaseClassifier): + """ + Классификатор породы для кошек и собак + Использует ResNet50 с дообучением на 37 пород + """ + + def __init__(self, device: str = 'cpu'): + """ + Args: + model_path: путь к обученной модели (опционально) + device: устройство для инференса ('cpu', 'mps', 'cuda') + """ + self.device = device + + # Создаем модель ResNet50 с предобученными весами на ImageNet + self.model = models.resnet50(pretrained=True) + + # Заменяем последний слой на 37 классов (породы) + num_features = self.model.fc.in_features + self.model.fc = nn.Linear(num_features, 37) + model_path = "../../models/unet_cats_dogs.pth" + # Загружаем обученные веса если есть + self.loaded = False + if model_path and Path(model_path).exists(): + try: + self.model.load_state_dict(torch.load(model_path, map_location=device)) + self.loaded = True + print(f"✅ Загружен классификатор из {model_path}") + except Exception as e: + print(f"⚠️ Ошибка загрузки модели из {model_path}: {e}") + print(" Используем неподготовленную модель (будет работать плохо)") + else: + print(f"⚠️ Классификатор не найден ({model_path}), используем неподготовленную модель ") + print(" Для точной работы обучите модель: python train_classifier.py") + + # Перемещаем на устройство + self.model.to(device) + self.model.eval() + + print(f"📊 Модель на устройстве: {device}") + + def predict(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, Any]: + """ + Предсказание породы по изображению + + Args: + image: изображение в формате numpy array или PIL Image + + Returns: + Словарь с результатами: + - breed_id: ID породы (1-37) + - breed_name: название породы + - species_id: ID вида (1-кошка, 2-собака) + - species_name: название вида + - confidence: уверенность (0-1) + - loaded: загружена ли модель + """ + # Проверка загрузки модели + if not self.loaded: + return { + "breed_id": 0, + "breed_name": "Unknown", + "species_id": 0, + "species_name": "Unknown", + "confidence": 0.0, + "loaded": False, + "error": "Model not loaded" + } + + try: + # Подготовка изображения + if isinstance(image, np.ndarray): + image = Image.fromarray(image) + + # Ресайз до 224x224 (стандартный вход ResNet) + image = image.resize((224, 224)) + image_array = np.array(image, dtype=np.float32) / 255.0 + image_array = image_array.transpose(2, 0, 1) # HWC -> CHW + + # Нормализация для ImageNet (как у ResNet) + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + for i in range(3): + image_array[i] = (image_array[i] - mean[i]) / std[i] + + # Создаем тензор + image_tensor = torch.FloatTensor(image_array).unsqueeze(0).to(self.device) + + # Инференс + with torch.no_grad(): + outputs = self.model(image_tensor) + probabilities = torch.softmax(outputs, dim=1) + confidence, predicted = torch.max(probabilities, 1) + + # Формируем результат + breed_id = predicted.item() + 1 # ID начинаются с 1 + confidence_score = confidence.item() + + # Определяем вид (кошка или собака) + # 1-25 кошки, 26-37 собаки + species_id = 1 if breed_id <= 25 else 2 + + return { + "breed_id": breed_id, + "breed_name": BREED_MAP.get(breed_id, "Unknown"), + "species_id": species_id, + "species_name": SPECIES_MAP.get(species_id, "Unknown"), + "confidence": confidence_score, + "loaded": True, + "error": None + } + + except Exception as e: + return { + "breed_id": 0, + "breed_name": "Error", + "species_id": 0, + "species_name": "Error", + "confidence": 0.0, + "loaded": self.loaded, + "error": str(e) + } + + def classify(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, Any]: + """Алиас для predict (для совместимости с BaseClassifier)""" + return self.predict(image) + + def get_info(self) -> Dict[str, Any]: + """Информация о классификаторе""" + return { + "name": self.__class__.__name__, + "type": "classification", + "target": "pet_breeds", + "num_classes": 37, + "loaded": self.loaded, + "device": self.device, + "model_architecture": "ResNet50", + "pretrained": True, + "breeds": list(BREED_MAP.values())[:5] + ["..."], # Показываем 5 пород + "total_breeds": len(BREED_MAP) + } \ No newline at end of file diff --git a/src/classifiers/simple_classifier.py b/src/classifiers/simple_classifier.py new file mode 100644 index 0000000..f629e47 --- /dev/null +++ b/src/classifiers/simple_classifier.py @@ -0,0 +1,40 @@ +from typing import Union, Dict, Any + +import numpy as np +from PIL import Image + +from src.classifiers.bone_condition_classifier import BoneConditionClassifier + + +# Упрощенная версия для быстрого старта +class SimpleBoneConditionClassifier(BoneConditionClassifier): + """ + Упрощенная версия классификатора для демонстрации + Использует предобученный ResNet50 без дообучения + """ + + def __init__(self, device: str = 'cpu'): + super().__init__( + model_path=None, # Не загружаем веса + device=device, + num_classes=2 # Только 2 класса: норма/аномалия + ) + print("⚠️ Используется упрощенная версия классификатора") + print(" Для точной работы обучите модель на медицинских данных") + + def predict(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, Any]: + """Упрощенное предсказание для демонстрации""" + if not self.loaded: + # Имитация работы для демо + return { + "condition_id": 0, + "condition_name": "Normal (Demo)", + "severity": "None", + "confidence": 0.95, + "loaded": True, + "error": None, + "is_abnormal": False, + "risk_level": "Low", + "demo_mode": True + } + return super().predict(image) \ No newline at end of file diff --git a/src/core/__init__.py b/src/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/core/config.py b/src/core/config.py new file mode 100644 index 0000000..3717318 --- /dev/null +++ b/src/core/config.py @@ -0,0 +1,26 @@ +from pydantic import BaseModel +from typing import Literal, Dict, Any + + +class PipelineConfig(BaseModel): + """Конфигурация пайплайна""" + + # Какой сегментатор использовать + segmentator: Literal["unet", "totalsegmentator", "custom"] = "unet" + + # Какой классификатор использовать + classifier: Literal["pet_breed", "bone_condition", "disease"] = "pet_breed" + + # Дополнительные параметры + params: Dict[str, Any] = { + "input_size": (256, 256), + "device": "auto", + "threshold": 0.5 + } + + # Режим работы + mode: Literal["development", "medical"] = "development" + + +# Глобальная конфигурация +config = PipelineConfig() \ No newline at end of file diff --git a/src/core/orchestrator.py b/src/core/orchestrator.py new file mode 100644 index 0000000..336f033 --- /dev/null +++ b/src/core/orchestrator.py @@ -0,0 +1,120 @@ +# src/core/orchestrator.py +import torch +import numpy as np +from PIL import Image +from typing import Dict, Any + +from src.classifiers.bone_condition_classifier import BoneConditionClassifier +from src.core.config import PipelineConfig +from src.core.registry import ModelRegistry +from src.classifiers.pet_breed_classifier import PetBreedClassifier +from src.quality.universal_scorer import UniversalQualityScorer +from src.segmentators.totalsegmentator_wrapper import TotalSegmentatorWrapper +from src.segmentators.unet_segmentator import UNetSegmentator + +# Регистрируем модели +ModelRegistry.register_segmentator("unet", UNetSegmentator) +ModelRegistry.register_segmentator("totalsegmentator", TotalSegmentatorWrapper) + +ModelRegistry.register_classifier("pet_breed", PetBreedClassifier) +ModelRegistry.register_classifier("bone_condition", BoneConditionClassifier) + + +class Orchestrator: + """Главный оркестратор для управления пайплайном""" + + def __init__(self, config: PipelineConfig = None): + self.config = config or PipelineConfig() + self.device = self._get_device() + # Инициализация компонентов + self._init_components() + + def _get_device(self): + """Определение устройства""" + if torch.backends.mps.is_available(): + return "mps" + elif torch.cuda.is_available(): + return "cuda" + else: + return "cpu" + + def _init_components(self): + """Инициализация компонентов на основе конфигурации""" + # Сегментатор + segmentator_class = ModelRegistry.get_segmentator(self.config.segmentator) + if segmentator_class: + self.segmentator = segmentator_class(device=self.device) + else: + raise ValueError(f"Сегментатор {self.config.segmentator} не найден") + + # Классификатор + classifier_class = ModelRegistry.get_classifier(self.config.classifier) + if classifier_class: + self.classifier = classifier_class(device=self.device) + else: + print(f"⚠️ Классификатор {self.config.classifier} не найден") + self.classifier = None + + # Quality Scorer (универсальный) + self.quality_scorer = UniversalQualityScorer() + + def analyze(self, image) -> Dict[str, Any]: + """Полный анализ изображения""" + # 1. Сегментация + mask = self.segmentator.segment(image) + + # 2. Классификация (если есть) + classification_result = None + if self.classifier: + classification_result = self.classifier.classify(image) + + # 3. Оценка качества + if isinstance(image, Image.Image): + image_array = np.array(image) + else: + image_array = image + + quality_result = self.quality_scorer.evaluate( + image=image_array, + segmentation=mask, + config=self.config + ) + + # 4. Сборка результата + return { + "mask": mask, + "classification": classification_result, + "quality": quality_result, + "pipeline_info": { + "segmentator": self.segmentator.get_info(), + "classifier": self.classifier.get_info() if self.classifier else None, + "mode": self.config.mode + } + } + + def switch_mode(self, mode: str): + """Переключение режима работы""" + self.config.mode = mode + + # Если переключаемся на медицинский режим + if mode == "medical": + self.config.segmentator = "totalsegmentator" + self.config.classifier = "bone_condition" + else: + self.config.segmentator = "unet" + self.config.classifier = "pet_breed" + + # Переинициализация компонентов + self._init_components() + print(f"✅ Переключено на режим: {mode}") + + def get_info(self) -> Dict[str, Any]: + """Информация о текущем пайплайне""" + return { + "mode": self.config.mode, + "segmentator": self.segmentator.get_info(), + "classifier": self.classifier.get_info() if self.classifier else None, + "device": self.device, + "available_segmentators": ModelRegistry.list_segmentators(), + "available_classifiers": ModelRegistry.list_classifiers() + } \ No newline at end of file diff --git a/src/core/registry.py b/src/core/registry.py new file mode 100644 index 0000000..8a48bd4 --- /dev/null +++ b/src/core/registry.py @@ -0,0 +1,40 @@ +from typing import Dict, Type +from src.segmentators.base import BaseSegmentator +from src.classifiers.base import BaseClassifier + + +class ModelRegistry: + """Регистр моделей для плагинной системы""" + + _segmentators: Dict[str, Type[BaseSegmentator]] = {} + _classifiers: Dict[str, Type[BaseClassifier]] = {} + + @classmethod + def register_segmentator(cls, name: str, segmentator_class: Type[BaseSegmentator]): + """Регистрация сегментатора""" + cls._segmentators[name] = segmentator_class + + @classmethod + def register_classifier(cls, name: str, classifier_class: Type[BaseClassifier]): + """Регистрация классификатора""" + cls._classifiers[name] = classifier_class + + @classmethod + def get_segmentator(cls, name: str): + """Получение сегментатора по имени""" + return cls._segmentators.get(name) + + @classmethod + def get_classifier(cls, name: str): + """Получение классификатора по имени""" + return cls._classifiers.get(name) + + @classmethod + def list_segmentators(cls): + """Список доступных сегментаторов""" + return list(cls._segmentators.keys()) + + @classmethod + def list_classifiers(cls): + """Список доступных классификаторов""" + return list(cls._classifiers.keys()) \ No newline at end of file diff --git a/src/data_loader.py b/src/data_loader.py index fb47004..84e6636 100644 --- a/src/data_loader.py +++ b/src/data_loader.py @@ -1,7 +1,5 @@ -# src/data_loader.py - легко заменить на DICOM import torch -from torch.utils.data import Dataset, DataLoader -from torchvision import transforms +from torch.utils.data import Dataset import numpy as np from PIL import Image import os diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..962552a --- /dev/null +++ b/src/main.py @@ -0,0 +1,49 @@ +from fastapi import FastAPI +from fastapi.staticfiles import StaticFiles +from pathlib import Path + +from starlette.responses import JSONResponse + +from src.api.endpoints import router + +# Создаем приложение +app = FastAPI( + title="Bone Quality Assessment API", + description="API для оценки качества медицинских изображений", + version="1.0.0", + docs_url="/docs", + redoc_url="/redoc", + openapi_url="/openapi.json" +) + +# Подключаем статические файлы +static_dir = Path(__file__).parent / "api/static" +if static_dir.exists(): + app.mount("/static", StaticFiles(directory=str(static_dir)), name="static") + print(f"✅ Статика подключена: {static_dir}") +else: + print(f"⚠️ Папка static не найдена: {static_dir}") + +app.include_router(router) + +# Глобальный обработчик 404 +@app.exception_handler(404) +async def not_found_handler(request, exc): + return JSONResponse( + status_code=404, + content={ + "error": "Not Found", + "message": f"Endpoint {request.url.path} not found", + "available_endpoints": [ + "/", + "/docs", + "/redoc", + "/openapi.json", + "/api/v1/health", + "/api/v1/analyze (POST)", + "/api/v1/switch_mode", + "/api/v1/info", + "/static/" + ] + } + ) \ No newline at end of file diff --git a/src/model/__init__.py b/src/model/__init__.py index 7179c4b..baf6e16 100644 --- a/src/model/__init__.py +++ b/src/model/__init__.py @@ -3,12 +3,10 @@ """ __all__ = [ - # U-Net "UNet", "DoubleConv", "DownBlock", "UpBlock", - "SegmentationEngine" ] diff --git a/src/quality/__init__.py b/src/quality/__init__.py index 4ada364..c29447c 100644 --- a/src/quality/__init__.py +++ b/src/quality/__init__.py @@ -2,7 +2,7 @@ Оценка качества медицинских изображений """ -from src.quality.quality_scorer_old import QualityScorer +# from src.quality.quality_scorer_old import QualityScorer from src.quality.artifact_detector import ArtifactDetector from src.quality.position_validator import PositionValidator @@ -16,7 +16,7 @@ QUALITY_ISSUES = { } __all__ = [ - "QualityScorer", + # "QualityScorer", "ArtifactDetector", "PositionValidator", "QUALITY_ISSUES", diff --git a/src/quality/medical_quality.py b/src/quality/medical_quality.py index 8431cbb..f36dcb7 100644 --- a/src/quality/medical_quality.py +++ b/src/quality/medical_quality.py @@ -3,8 +3,7 @@ from typing import Dict, Any import numpy as np from monai.transforms.utils import ndimage -from src import QualityScorer - +from src.quality.quality_scorer import QualityScorer class MedicalQualityScorer(QualityScorer): def __init__(self, anatomy_type="spine"): diff --git a/src/quality/quality_scorer_old.py b/src/quality/quality_scorer_old.py deleted file mode 100644 index 07a5ef7..0000000 --- a/src/quality/quality_scorer_old.py +++ /dev/null @@ -1,249 +0,0 @@ -import numpy as np -from scipy import ndimage -from typing import Dict, Any, Optional, Tuple - - -class QualityScorer: - """Оценка качества сегментации (работает с любыми объектами)""" - - def __init__(self): - self.quality_metrics = {} - - def check_artifact(self, segmentation: np.ndarray, context: str = "single_object") -> Dict[str, Any]: - """ - Проверка артефактов с учетом контекста - - Args: - segmentation: бинарная маска - context: "single_object" (кот/собака) или "multiple_objects" (позвонки) - """ - if segmentation is None or segmentation.sum() == 0: - return {"artifact": True, "type": "empty", "severity": "critical"} - - # Находим все объекты - labeled, num_features = ndimage.label(segmentation) - - # Разная логика для разных контекстов - if context == "multiple_objects": - # Для позвонков: ожидаем 3-5 объектов - if num_features < 3: - return {"artifact": True, "type": "too_few_objects", "severity": "high"} - elif num_features > 7: - return {"artifact": True, "type": "too_many_objects", "severity": "medium"} - else: - # Для кошек/собак (или бедренной кости): ожидаем 1 объект - if num_features > 1: - # Проверяем размер объектов - sizes = [np.sum(labeled == i) for i in range(1, num_features + 1)] - main_object_size = max(sizes) - total_size = sum(sizes) - - # Если один объект намного больше остальных - это шум - if main_object_size / total_size > 0.8: - return {"artifact": False, "num_objects": num_features, - "minor_objects": num_features - 1} - else: - return {"artifact": True, "type": "fragmented", "severity": "high"} - - return { - "artifact": False, - "num_objects": num_features, - "object_sizes": [np.sum(labeled == i) for i in range(1, num_features + 1)] - } - - def check_confidence(self, segmentation: np.ndarray, probability_map: np.ndarray) -> Dict[str, Any]: - """ - Проверка уверенности модели в сегментации - """ - # Если модель не уверена - это может быть артефакт - mean_confidence = probability_map[segmentation > 0].mean() - - return { - "confidence": float(mean_confidence), - "valid": mean_confidence > 0.5, # порог уверенности - "low_confidence_regions": (probability_map < 0.3).sum() / segmentation.sum() - } - - def check_position(self, segmentation: np.ndarray, - expected_center: Tuple[float, float] = (0.5, 0.5), - allowed_deviation: float = 0.25) -> Dict[str, Any]: - """ - Проверка положения объекта в кадре - - Args: - segmentation: бинарная маска сегментации [H, W] - expected_center: ожидаемый центр (y, x) в нормализованных координатах - allowed_deviation: максимальное допустимое отклонение - - Returns: - Словарь с результатами проверки - """ - if segmentation is None or segmentation.sum() == 0: - return {"position": None, "valid": False, "error": "empty_mask"} - - # Центр масс - y, x = np.where(segmentation > 0) - if len(y) == 0: - return {"position": None, "valid": False, "error": "no_pixels"} - - center_y, center_x = np.mean(y), np.mean(x) - h, w = segmentation.shape - - # Нормализованные координаты (0-1) - norm_y = center_y / h - norm_x = center_x / w - - # Отклонение от центра - deviation = np.sqrt((norm_y - expected_center[0]) ** 2 + - (norm_x - expected_center[1]) ** 2) - - # Объект должен быть в центре - valid = deviation < allowed_deviation - - # Проверка размера (не слишком маленький) - object_size = len(y) / (h * w) - if object_size < 0.005: # меньше 0.5% площади - valid = False - - return { - "position": (float(norm_y), float(norm_x)), - "valid": valid, - "deviation": float(deviation), - "object_size": float(object_size), - "center_pixels": (int(center_y), int(center_x)), - "expected_center": expected_center, - "allowed_deviation": allowed_deviation - } - - def check_contrast(self, image: np.ndarray, segmentation: np.ndarray) -> Dict[str, Any]: - """ - Проверка контрастности области объекта - - Args: - image: исходное изображение [H, W] или [H, W, C] - segmentation: бинарная маска сегментации [H, W] - - Returns: - Словарь с результатами проверки - """ - if segmentation is None or segmentation.sum() == 0: - return {"valid": False, "contrast": 0, "error": "empty_mask"} - - # Если изображение цветное, конвертируем в grayscale - if len(image.shape) == 3: - image_gray = np.mean(image, axis=2) - else: - image_gray = image - - # Значения пикселей внутри объекта - object_pixels = image_gray[segmentation > 0] - if len(object_pixels) == 0: - return {"valid": False, "contrast": 0, "error": "no_pixels"} - - # Значения пикселей снаружи объекта (фон) - background_pixels = image_gray[segmentation == 0] - - # Контраст: отношение std/mean внутри объекта - contrast = np.std(object_pixels) / (np.mean(object_pixels) + 1e-8) - - # Отношение средних (объект/фон) - mean_ratio = np.mean(object_pixels) / (np.mean(background_pixels) + 1e-8) - - return { - "valid": contrast > 0.08, # порог - "contrast": float(contrast), - "mean_ratio": float(mean_ratio), - "mean_intensity": float(np.mean(object_pixels)), - "std_intensity": float(np.std(object_pixels)), - "background_mean": float(np.mean(background_pixels)) - } - - def evaluate(self, image: Optional[np.ndarray] = None, - segmentation: Optional[np.ndarray] = None) -> Dict[str, Any]: - """ - Полная оценка качества - - Args: - image: исходное изображение [H, W] или [H, W, C] (опционально) - segmentation: бинарная маска сегментации [H, W] - - Returns: - Словарь с полной оценкой качества - """ - results = { - "overall_quality": "GOOD", - "severity": "LOW", - "issues": [], - "metrics": {}, - "confidence": 0.85, - "position": {}, - "artifact": {} - } - - # 1. Проверка артефактов - artifact_results = self.check_artifact(segmentation) - results["metrics"]["artifact"] = artifact_results - - if artifact_results.get("artifact", False): - results["issues"].append({ - "type": "ARTIFACT", - "details": artifact_results.get("type", "unknown"), - "severity": artifact_results.get("severity", "MEDIUM") - }) - results["severity"] = max(results["severity"], - artifact_results.get("severity", "MEDIUM")) - - # 2. Проверка позиции - position_results = self.check_position(segmentation) - results["metrics"]["position"] = position_results - - if not position_results.get("valid", False): - results["issues"].append({ - "type": "POSITION_ERROR", - "details": f"Object at {position_results.get('position')}, deviation {position_results.get('deviation', 0):.3f}", - "severity": "HIGH" - }) - results["severity"] = "HIGH" - - # 3. Проверка контраста (если есть изображение) - if image is not None: - contrast_results = self.check_contrast(image, segmentation) - results["metrics"]["contrast"] = contrast_results - - if not contrast_results.get("valid", False): - results["issues"].append({ - "type": "LOW_CONTRAST", - "details": f"Contrast: {contrast_results.get('contrast', 0):.3f}", - "severity": "MEDIUM" - }) - if results["severity"] != "HIGH": - results["severity"] = "MEDIUM" - - # Интегральная оценка - if not results["issues"]: - results["overall_quality"] = "GOOD" - results["confidence"] = 0.9 - elif results["severity"] == "HIGH": - results["overall_quality"] = "POOR" - results["confidence"] = 0.5 - else: - results["overall_quality"] = "WARNING" - results["confidence"] = 0.7 - - return results - - -# # В медицинском сервисе -# segmentation = model.segment(dicom_image) -# spine_check = check_spine_segmentation(segmentation) -# -# if spine_check["valid"]: -# print("✅ Позвоночник сегментирован правильно") -# print(f" Найдено позвонков: {spine_check['metrics']['num_vertebrae']}") -# print(f" Выравнивание: {spine_check['metrics']['alignment_deviation']:.1f} пикселей") -# else: -# print("❌ Обнаружены проблемы с сегментацией позвоночника:") -# for issue in spine_check["issues"]: -# print(f" • {issue['type']}: {issue['details']}") -# for rec in spine_check["recommendations"]: -# print(f" 💡 {rec}") \ No newline at end of file diff --git a/src/quality/universal_scorer.py b/src/quality/universal_scorer.py new file mode 100644 index 0000000..d8d2368 --- /dev/null +++ b/src/quality/universal_scorer.py @@ -0,0 +1,7 @@ +from src.quality.quality_scorer import QualityScorer + + +class UniversalQualityScorer(QualityScorer): + + def __init__(self): + super().__init__() \ No newline at end of file diff --git a/src/run.py b/src/run.py new file mode 100644 index 0000000..b066fa4 --- /dev/null +++ b/src/run.py @@ -0,0 +1,12 @@ +import uvicorn + +if __name__ == "__main__": + print("🚀 Запуск Bone Quality Assessment API") + print("🌐 http://localhost:8000") + uvicorn.run( + "src.main:app", + host="0.0.0.0", + port=8000, + reload=True, + log_level="info" + ) \ No newline at end of file diff --git a/src/segmentators/__init__.py b/src/segmentators/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/segmentators/base.py b/src/segmentators/base.py new file mode 100644 index 0000000..3090aac --- /dev/null +++ b/src/segmentators/base.py @@ -0,0 +1,29 @@ +from abc import ABC, abstractmethod +import numpy as np +from PIL import Image +from typing import Union, Dict, Any + + +class BaseSegmentator(ABC): + """Базовый класс для всех сегментаторов""" + + @abstractmethod + def segment(self, image: Union[np.ndarray, Image.Image]) -> np.ndarray: + """ + Сегментация изображения + + Args: + image: входное изображение + + Returns: + Бинарная маска [H, W] где 1 = объект + """ + pass + + @abstractmethod + def get_info(self) -> Dict[str, Any]: + """Информация о модели""" + return { + "name": self.__class__.__name__, + "type": "segmentation" + } \ No newline at end of file diff --git a/src/segmentators/totalsegmentator_wrapper.py b/src/segmentators/totalsegmentator_wrapper.py new file mode 100644 index 0000000..c123698 --- /dev/null +++ b/src/segmentators/totalsegmentator_wrapper.py @@ -0,0 +1,35 @@ +from src.segmentators.base import BaseSegmentator +import numpy as np + + +class TotalSegmentatorWrapper(BaseSegmentator): + """Обертка для TotalSegmentator (для костей)""" + + def __init__(self): + try: + from totalsegmentator import TotalSegmentator + self.model = TotalSegmentator() + self.loaded = True + print("✅ TotalSegmentator загружен") + except ImportError: + print("⚠️ TotalSegmentator не установлен") + self.loaded = False + + def segment(self, image): + if not self.loaded: + return np.zeros((512, 512), dtype=np.int64) + + # TotalSegmentator ожидает DICOM или NIfTI + # Здесь адаптация под входной формат + # Реальная реализация зависит от TotalSegmentator API + + # Пока возвращаем заглушку + return np.ones((512, 512), dtype=np.int64) + + def get_info(self): + return { + "name": "TotalSegmentator", + "type": "segmentation", + "target": "bones", + "loaded": self.loaded + } \ No newline at end of file diff --git a/src/segmentators/unet_segmentator.py b/src/segmentators/unet_segmentator.py new file mode 100644 index 0000000..04433e5 --- /dev/null +++ b/src/segmentators/unet_segmentator.py @@ -0,0 +1,61 @@ +# src/segmentators/unet_segmentator.py +from pathlib import Path +import torch +import numpy as np +from PIL import Image + +from src import UNet +from src.segmentators.base import BaseSegmentator + + +class UNetSegmentator(BaseSegmentator): + """Сегментатор на основе U-Net (для кошек/собак)""" + + def __init__(self, model_path="models/unet_cats_dogs.pth", device='cpu'): + self.device = device + self.model = UNet(in_channels=3, out_classes=2) + + if Path(model_path).exists(): + self.model.load_state_dict(torch.load(model_path, map_location=device)) + self.model.to(device) + self.model.eval() + self.loaded = True + print(f"✅ U-Net загружен из {model_path}") + else: + print(f"⚠️ U-Net не найден: {model_path}") + self.loaded = False + + def segment(self, image): + if not self.loaded: + return np.zeros((256, 256), dtype=np.int64) + + # Подготовка изображения + if isinstance(image, np.ndarray): + image = Image.fromarray(image) + + original_size = image.size + image_resized = image.resize((256, 256)) + image_array = np.array(image_resized, dtype=np.float32) / 255.0 + image_array = image_array.transpose(2, 0, 1) + image_tensor = torch.FloatTensor(image_array).unsqueeze(0).to(self.device) + + # Инференс + with torch.no_grad(): + output = self.model(image_tensor) + probs = torch.softmax(output, dim=1) + mask = probs.argmax(dim=1).squeeze().cpu().numpy() + + # Ресайз к оригиналу + mask_pil = Image.fromarray(mask.astype(np.uint8)) + mask_pil = mask_pil.resize(original_size, resample=Image.NEAREST) + + return np.array(mask_pil) + + def get_info(self): + return { + "name": "U-Net", + "type": "segmentation", + "target": "cats_and_dogs", + "classes": 2, + "loaded": self.loaded + } \ No newline at end of file diff --git a/src/utils/utils.py b/src/utils/utils.py new file mode 100644 index 0000000..8c0a8cd --- /dev/null +++ b/src/utils/utils.py @@ -0,0 +1,11 @@ +import torch + + +def get_device(): + """Определение устройства""" + if torch.backends.mps.is_available(): + return torch.device("mps") + elif torch.cuda.is_available(): + return torch.device("cuda") + else: + return torch.device("cpu") \ No newline at end of file