Update api_app.py

api_app.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 import io
 import numpy as np
 import os
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional # Dodano Optional
 import logging # Dodajemy import modułu logging
 import cv2 # Dodajemy OpenCV
 import base64 # Dodajemy base64
@@ -17,6 +17,7 @@ import base64 # Dodajemy base64
 from pytorch_grad_cam import GradCAMPlusPlus
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from huggingface_hub import hf_hub_download # Do pobierania modelu z Huba
+from pydantic import BaseModel # Dodano Pydantic
 
 # --- Konfiguracja Logowania ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -31,6 +32,14 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 IMAGENET_MEAN = [0.485, 0.456, 0.406]
 IMAGENET_STD = [0.229, 0.224, 0.225]
 
+# --- Konfigurowalne progi (ze zmiennych środowiskowych z wartościami domyślnymi) ---
+COLORFULNESS_THRESHOLD = float(os.environ.get("COLORFULNESS_THRESHOLD", 15))
+UNIFORMITY_THRESHOLD = float(os.environ.get("UNIFORMITY_THRESHOLD", 10))
+ASPECT_RATIO_MIN = float(os.environ.get("ASPECT_RATIO_MIN", 0.4))
+ASPECT_RATIO_MAX = float(os.environ.get("ASPECT_RATIO_MAX", 2.5))
+MC_DROPOUT_SAMPLES = int(os.environ.get("MC_DROPOUT_SAMPLES", 25))
+UNCERTAINTY_THRESHOLD_STD = float(os.environ.get("UNCERTAINTY_THRESHOLD_STD", 0.08))
+
 # Globalne zmienne dla modelu i transformacji
 model_instance = None
 transform_pipeline = None
@@ -84,10 +93,57 @@ def initialize_model():
     transform_pipeline = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD) # Upewnij się, że IMAGENET_MEAN i IMAGENET_STD są zdefiniowane
     ])
     logger.info(f"Model BI-RADS classifier initialized successfully on device: {DEVICE}")
 
+# --- Funkcja do predykcji z kwantyfikacją niepewności (MC Dropout) ---
+def predict_with_mc_dropout(current_model_instance, input_tensor_on_device): # Usunięto domyślne wartości, będą brane z globalnych zmiennych
+    """
+    Wykonuje predykcję z użyciem Monte Carlo Dropout do oszacowania niepewności.
+    """
+    logger.info(f"Performing MC Dropout with {MC_DROPOUT_SAMPLES} samples. Uncertainty threshold (std): {UNCERTAINTY_THRESHOLD_STD}")
+
+    original_mode_is_training = current_model_instance.training
+    current_model_instance.train()  # Włącz warstwy dropout
+
+    all_probs_list = []
+    with torch.no_grad(): # Gradienty nie są potrzebne do samego przejścia w przód
+        for _ in range(n_samples):
+            output = current_model_instance(input_tensor_on_device) # Użyj MC_DROPOUT_SAMPLES
+            probs_tensor = torch.nn.functional.softmax(output, dim=1)
+            all_probs_list.append(probs_tensor.cpu().numpy())
+
+    # Przywróć oryginalny tryb modelu
+    if not original_mode_is_training:
+        current_model_instance.eval()
+
+    predictions_stack = np.vstack(all_probs_list) # Kształt: (n_samples, num_classes)
+    mean_probabilities = np.mean(predictions_stack, axis=0) # Kształt: (num_classes,)
+    std_dev_probabilities = np.std(predictions_stack, axis=0) # Kształt: (num_classes,)
+
+    predicted_class_index = np.argmax(mean_probabilities)
+    confidence_in_predicted_class = mean_probabilities[predicted_class_index]
+    
+    # Użyj średniej odchyleń standardowych prawdopodobieństw wszystkich klas jako metryki niepewności
+    uncertainty_metric = np.mean(std_dev_probabilities) 
+
+    is_uncertain = uncertainty_metric > UNCERTAINTY_THRESHOLD_STD # Użyj UNCERTAINTY_THRESHOLD_STD
+    logger.info(f"MC Dropout Results: Predicted Index: {predicted_class_index}, Confidence: {confidence_in_predicted_class:.4f}, Uncertainty (avg_std): {uncertainty_metric:.4f}, Is Uncertain: {is_uncertain}")
+
+    birads_category_if_confident = predicted_class_index + 1
+    base_result = {
+        "birads": birads_category_if_confident if not is_uncertain else None,
+        "confidence": float(confidence_in_predicted_class) if not is_uncertain else None,
+        "interpretation": interpretations_dict.get(birads_category_if_confident, "Nieznana klasyfikacja") if not is_uncertain \
+                          else f"Model jest niepewny co do tego obrazu (niepewność: {uncertainty_metric:.4f}). Sprawdź jakość i typ obrazu. Może to być obraz spoza domeny medycznej.",
+        "class_probabilities": {str(j + 1): float(mean_probabilities[j]) for j in range(len(mean_probabilities))},
+        "grad_cam_image_base64": None, # Zostanie wypełnione później, jeśli pewne
+        "error": "High prediction uncertainty" if is_uncertain else None,
+        "details": f"Uncertainty metric ({uncertainty_metric:.4f}) {'przekroczyła' if is_uncertain else 'jest w granicach'} progu ({UNCERTAINTY_THRESHOLD_STD})."
+    }
+    return base_result, predicted_class_index # Zwróć również indeks dla Grad-CAM
+
 # --- Funkcja do tworzenia obrazu z nało��oną mapą Grad-CAM (zaadaptowana z app.py) ---
 def create_grad_cam_overlay_image(original_pil_image: Image.Image, grayscale_cam: np.ndarray, birads_category: int, transparency: float = 0.5) -> Image.Image:
     """Tworzy obraz PIL z nałożoną mapą Grad-CAM."""
@@ -127,10 +183,60 @@ def create_grad_cam_overlay_image(original_pil_image: Image.Image, grayscale_cam
         logger.info("Obraz Grad-CAM overlay pomyślnie utworzony.")
         return Image.fromarray(final_image_np)
     except Exception as e:
-        logger.error(f"Błąd podczas tworzenia obrazu Grad-CAM overlay: {e}", exc_info=True)
+        logger.error(f"Błąd podczas tworzenia obrazu Grad-CAM overlay: {e}", exc_info=True) 
         return None
 
+# --- Funkcja do heurystycznych testów OOD ---
+def run_heuristic_ood_checks(pil_image: Image.Image, request_id: str) -> bool:
+    """
+    Wykonuje zestaw prostych heurystyk do wykrywania obrazów spoza dystrybucji.
+    Zwraca True, jeśli obraz przeszedł testy (prawdopodobnie jest OK), False jeśli nie.
+    """
+    logger.info(f"[RequestID: {request_id}] Uruchamianie heurystycznych testów OOD...")
+    width, height = pil_image.size
+
+    # Heurystyka 1: Sprawdzenie proporcji obrazu
+    aspect_ratio = width / height
+    if not (ASPECT_RATIO_MIN < aspect_ratio < ASPECT_RATIO_MAX): # Użyj zmiennych konfiguracyjnych
+        logger.warning(f"[RequestID: {request_id}] Heurystyka OOD: Nietypowe proporcje obrazu: {aspect_ratio:.2f}. Odrzucam.")
+        return False
+
+    # Heurystyka 2: Analiza "kolorowości" (dla obrazów RGB)
+    # Mammografie są w skali szarości; jeśli obraz jest kolorowy, ta metryka będzie wysoka.
+    # Zakładamy, że pil_image może być już w trybie RGB lub zostanie do niego skonwertowany.
+    # Jeśli pil_image jest w trybie 'L', można by tę heurystykę pominąć lub dostosować.
+    # Dla obrazu w skali szarości skonwertowanego do RGB, R=G=B, więc std_per_pixel_across_channels będzie bliskie 0.
+    img_rgb_for_color_check = pil_image.convert('RGB') # Upewnijmy się, że pracujemy na RGB dla tej heurystyki
+    img_np_rgb = np.array(img_rgb_for_color_check)
+    std_per_pixel_across_channels = np.std(img_np_rgb, axis=2) # Odch. std. dla każdego piksela po kanałach R,G,B
+    mean_std_across_channels = np.mean(std_per_pixel_across_channels)
+    if mean_std_across_channels > COLORFULNESS_THRESHOLD: # Użyj zmiennej konfiguracyjnej
+        logger.warning(f"[RequestID: {request_id}] Heurystyka OOD: Obraz wydaje się zbyt kolorowy. Średnie odch. std. między kanałami: {mean_std_across_channels:.2f}. Odrzucam.")
+        return False
+
+    # Heurystyka 3: Sprawdzenie, czy obraz nie jest prawie jednolity (dominująca jasność/ciemność)
+    # Konwertujemy do skali szarości dla tej analizy
+    gray_image = pil_image.convert('L')
+    std_dev_intensity = np.std(np.array(gray_image))
+    if std_dev_intensity < UNIFORMITY_THRESHOLD: # Użyj zmiennej konfiguracyjnej
+        logger.warning(f"[RequestID: {request_id}] Heurystyka OOD: Obraz wydaje się zbyt jednolity (mało zróżnicowania jasności). Odch. std. intensywności: {std_dev_intensity:.2f}. Odrzucam.")
+        return False
+    
+    logger.info(f"[RequestID: {request_id}] Heurystyczne testy OOD zakończone pomyślnie. Obraz wygląda na potencjalnie poprawny.")
+    return True
+
 # --- Aplikacja FastAPI ---
+# --- Definicja modelu odpowiedzi Pydantic ---
+class PredictionResult(BaseModel):
+    birads: Optional[int] = None
+    confidence: Optional[float] = None
+    interpretation: str
+    class_probabilities: Dict[str, float]
+    grad_cam_image_base64: Optional[str] = None
+    error: Optional[str] = None
+    details: Optional[str] = None
+
+
 app = FastAPI(title="BI-RADS Mammography Classification API")
 
 @app.on_event("startup")
@@ -139,7 +245,7 @@ async def startup_event():
     logger.info("Rozpoczynanie eventu startup aplikacji FastAPI.")
     initialize_model()
 
-@app.post("/predict/", response_model=List[Dict[str, Any]])
+@app.post("/predict/", response_model=List[PredictionResult]) # Użycie modelu Pydantic
 async def predict_image(file: UploadFile = File(...)):
     """
     Endpoint do klasyfikacji obrazu mammograficznego.
@@ -156,79 +262,93 @@ async def predict_image(file: UploadFile = File(...)):
     try:
         logger.info(f"[RequestID: {request_id}] Odczytywanie i przetwarzanie wgranego pliku...")
         contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        logger.info(f"[RequestID: {request_id}] Plik obrazu pomyślnie odczytany i przekonwertowany do RGB.")
+        # Wczytaj obraz, ale jeszcze nie konwertuj do RGB, aby heurystyki mogły działać na bardziej "surowych" danych
+        image_pil_original = Image.open(io.BytesIO(contents))
+        logger.info(f"[RequestID: {request_id}] Plik obrazu pomyślnie odczytany. Oryginalny tryb: {image_pil_original.mode}, Rozmiar: {image_pil_original.size}")
+
+        # --- Etap: Heurystyczne testy OOD (PRE-FILTR) ---
+        # Używamy .copy(), aby uniknąć potencjalnych modyfikacji oryginalnego obiektu image_pil_original
+        # przez funkcję run_heuristic_ood_checks, jeśli by takie wykonywała (np. konwersje inplace).
+        if not run_heuristic_ood_checks(image_pil_original.copy(), request_id):
+            # Heurystyki wykryły problem
+            return JSONResponse(
+                status_code=400, # Bad Request
+                content=[{ # API oczekuje listy wyników
+                    "error": "Image does not appear to be a valid medical mammogram based on initial checks.",
+                    "details": "Heurystyczne testy OOD nie powiodły się. Obraz odrzucony przed analizą przez model AI."
+                }]
+            )
+        image = image_pil_original.convert("RGB") # Teraz konwertuj do RGB dla modelu głównego
     except Exception as e:
         logger.error(f"[RequestID: {request_id}] Błąd podczas odczytu pliku obrazu: {e}", exc_info=True)
         raise HTTPException(status_code=400, detail=f"Nie można odczytać pliku obrazu: {e}")
-
-    # Preprocessing
+    
+    # --- Preprocessing ---
     logger.info(f"[RequestID: {request_id}] Rozpoczynanie preprocessingu obrazu...")
     input_tensor = transform_pipeline(image).unsqueeze(0).to(DEVICE)
     logger.info(f"[RequestID: {request_id}] Preprocessing zakończony. Kształt tensora wejściowego: {input_tensor.shape}")
 
-    # Inference
-    logger.info(f"[RequestID: {request_id}] Rozpoczynanie inferencji modelu...")
-    with torch.no_grad():
-        model_outputs = model_instance(input_tensor)
-    logger.info(f"[RequestID: {request_id}] Inferencja zakończona.")
-    # Postprocessing
-    probs = torch.nn.functional.softmax(model_outputs, dim=1)
-    confidences, predicted_indices = torch.max(probs, 1)
-    
+    # --- Etap: Predykcja z kwantyfikacją niepewności (MC Dropout) ---
+    # Upewnij się, że model_instance, DEVICE, interpretations_dict są dostępne
+    mc_output_dict, predicted_idx_from_mc = predict_with_mc_dropout(
+        model_instance,
+        input_tensor,
+        # n_samples i uncertainty_threshold_std są teraz brane z globalnych zmiennych
+    )
+
     results = []
-    for i in range(len(predicted_indices)): # Pętla na wypadek przyszłego batch processingu
-        birads_category = predicted_indices[i].item() + 1
-        logger.info(f"[RequestID: {request_id}] Przewidziana kategoria BI-RADS: {birads_category}, Pewność: {confidences[i].item():.4f}")
-        confidence = confidences[i].item()
-        interpretation = interpretations_dict.get(birads_category, "Nieznana klasyfikacja")
-        
-        all_class_probs_tensor = probs[i].cpu().numpy()
-        class_probabilities = {str(j+1): float(all_class_probs_tensor[j]) for j in range(len(all_class_probs_tensor))}
-
-        # Generowanie Grad-CAM
-        grad_cam_map_serialized = None
-        grad_cam_image_base64 = None
-        logger.info(f"[RequestID: {request_id}] Rozpoczynanie generowania Grad-CAM dla kategorii {birads_category}...")
-        try:
-            # model_instance is already in eval mode from initialize_model()
-            # pytorch-grad-cam typically handles necessary gradient contexts.
-            target_layers = [model_instance.layer4[-1]] # Dla ResNet-18
-            cam_algorithm = GradCAMPlusPlus(model=model_instance, target_layers=target_layers)
-            
-            current_input_tensor_for_cam = input_tensor[i].unsqueeze(0).clone().detach().requires_grad_(True)
-            targets_for_cam = [ClassifierOutputTarget(predicted_indices[i].item())]
-            
-            grayscale_cam = cam_algorithm(input_tensor=current_input_tensor_for_cam, targets=targets_for_cam)
-            
-            if grayscale_cam is not None:
-                grad_cam_map_np = grayscale_cam[0, :]
-                logger.info(f"[RequestID: {request_id}] Grad-CAM wygenerowany pomyślnie.")
-
-                # Tworzenie obrazu z nałożoną mapą Grad-CAM
-                overlay_image_pil = create_grad_cam_overlay_image(original_pil_image=image, # oryginalny obraz PIL
-                                                                  grayscale_cam=grad_cam_map_np,
-                                                                  birads_category=birads_category)
-                if overlay_image_pil:
-                    buffered = io.BytesIO()
-                    overlay_image_pil.save(buffered, format="PNG") # Zapisz jako PNG
-                    grad_cam_image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-                    logger.info(f"[RequestID: {request_id}] Obraz Grad-CAM overlay zakodowany do base64.")
-                else:
-                    logger.warning(f"[RequestID: {request_id}] Nie udało się utworzyć obrazu Grad-CAM overlay.")
+    if mc_output_dict.get("error") == "High prediction uncertainty":
+        logger.warning(f"[RequestID: {request_id}] Wysoka niepewność predykcji: {mc_output_dict.get('details')}")
+        results.append(mc_output_dict) # Dodaj wynik z informacją o niepewności
+        return JSONResponse(content=results) # API oczekuje listy
+
+    # --- Jeśli predykcja jest pewna, kontynuuj z Grad-CAM ---
+    # mc_output_dict zawiera już 'birads', 'confidence', 'interpretation', 'class_probabilities'
+    birads_category_for_cam = mc_output_dict["birads"] # To jest już BI-RADS 1-5
+    # predicted_idx_from_mc to indeks 0-4
+
+    # Upewnij się, że model jest w trybie ewaluacji dla Grad-CAM
+    # Funkcja MC Dropout powinna go przywrócić, ale dla pewności:
+    model_instance.eval()
+
+    # Generowanie Grad-CAM
+    grad_cam_map_serialized = None # Zmienna zdefiniowana przed blokiem try
+    grad_cam_image_base64 = None   # Zmienna zdefiniowana przed blokiem try
+    # Poprawka: Użyj birads_category_for_cam zamiast niezdefiniowanej birads_category
+    logger.info(f"[RequestID: {request_id}] Rozpoczynanie generowania Grad-CAM dla kategorii {birads_category_for_cam}...")
+    try:
+        # model_instance is already in eval mode from initialize_model()
+        target_layers = [model_instance.layer4[-1]]
+        cam_algorithm = GradCAMPlusPlus(model=model_instance, target_layers=target_layers)
+
+        # Tensor wejściowy dla CAM musi mieć requires_grad=True
+        current_input_tensor_for_cam = input_tensor.clone().detach().requires_grad_(True)
+        targets_for_cam = [ClassifierOutputTarget(predicted_idx_from_mc)] # Użyj indeksu z MC Dropout
+
+        grayscale_cam = cam_algorithm(input_tensor=current_input_tensor_for_cam, targets=targets_for_cam)
+        if grayscale_cam is not None:
+            grad_cam_map_np = grayscale_cam[0, :]
+            logger.info(f"[RequestID: {request_id}] Grad-CAM wygenerowany pomyślnie.")
+
+            # Tworzenie obrazu z nałożoną mapą Grad-CAM
+            overlay_image_pil = create_grad_cam_overlay_image(original_pil_image=image, # oryginalny obraz PIL
+                                                              grayscale_cam=grad_cam_map_np,
+                                                              birads_category=birads_category_for_cam)
+            if overlay_image_pil:
+                buffered = io.BytesIO()
+                overlay_image_pil.save(buffered, format="PNG") # Zapisz jako PNG
+                grad_cam_image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+                logger.info(f"[RequestID: {request_id}] Obraz Grad-CAM overlay zakodowany do base64.")
             else:
-                logger.warning(f"[RequestID: {request_id}] Wygenerowany Grad-CAM jest None.")
-        except Exception as e:
-            logger.error(f"[RequestID: {request_id}] Błąd podczas generowania Grad-CAM w API: {e}", exc_info=True)
-
-        results.append({
-            "birads": birads_category,
-            "confidence": confidence,
-            "interpretation": interpretation,
-            "class_probabilities": class_probabilities,
-            "grad_cam_image_base64": grad_cam_image_base64 # Dodajemy obraz base64
-        })
-    
+                logger.warning(f"[RequestID: {request_id}] Nie udało się utworzyć obrazu Grad-CAM overlay.")
+        else:
+            logger.warning(f"[RequestID: {request_id}] Wygenerowany Grad-CAM jest None.")
+    except Exception as e:
+        logger.error(f"[RequestID: {request_id}] Błąd podczas generowania Grad-CAM w API: {e}", exc_info=True)
+        # grad_cam_image_base64 pozostanie None
+
+    mc_output_dict["grad_cam_image_base64"] = grad_cam_image_base64 # Zaktualizuj słownik wynikowy
+    results.append(mc_output_dict) # Dodaj finalny wynik (pewny, z Grad-CAM lub bez)
     logger.info(f"[RequestID: {request_id}] Przetwarzanie żądania /predict/ zakończone. Zwracam wyniki.")
     return JSONResponse(content=results)