Update api_app.py

api_app.py

@@ -1,153 +1,161 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
-import torch
-import torchvision.models as models
-import torchvision.transforms as transforms
-from PIL import Image
-import torch.nn as nn
-import io
-import numpy as np
-import os
-from typing import List, Dict, Any
-
-# Importy dla Grad-CAM
-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
-
-# --- Konfiguracja ---
-# Upewnij się, że te wartości są zgodne z Twoim repozytorium modelu
-HF_MODEL_REPO_ID = "Enterwar99/MODEL_MAMMOGRAFII"
-MODEL_FILENAME = "best_model.pth" # Nazwa pliku modelu w repozytorium
-
-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]
-
-# Globalne zmienne dla modelu i transformacji
-model_instance = None
-transform_pipeline = None
-
-interpretations_dict = {
-    1: "Wynik negatywny - brak zmian nowotworowych",
-    2: "Zmiana łagodna",
-    3: "Prawdopodobnie zmiana łagodna - zalecana kontrola",
-    4: "Podejrzenie zmiany złośliwej - zalecana biopsja",
-    5: "Wysoka podejrzliwość złośliwości - wymagana biopsja"
-}
-
-# --- Inicjalizacja modelu ---
-def initialize_model():
-    global model_instance, transform_pipeline
-    if model_instance is not None:
-        return
-
-    print(f"Pobieranie modelu {MODEL_FILENAME} z repozytorium {HF_MODEL_REPO_ID}...")
-    try:
-        model_pt_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=MODEL_FILENAME)
-    except Exception as e:
-        print(f"Błąd podczas pobierania modelu z Hugging Face Hub: {e}")
-        raise RuntimeError(f"Nie można pobrać modelu: {e}")
-
-    print(f"Inicjalizacja architektury modelu ResNet-18...")
-    model_arch = models.resnet18(weights=None)
-    num_feats = model_arch.fc.in_features
-    model_arch.fc = nn.Sequential(
-        nn.Dropout(0.5),
-        nn.Linear(num_feats, 5)
-    )
-    
-    print(f"Ładowanie wag modelu z {model_pt_path}...")
-    model_arch.load_state_dict(torch.load(model_pt_path, map_location=DEVICE))
-    model_arch.to(DEVICE)
-    model_arch.eval()
-    model_instance = model_arch
-
-    transform_pipeline = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
-    ])
-    print(f"Model BI-RADS classifier initialized successfully on device: {DEVICE}")
-
-# --- Aplikacja FastAPI ---
-app = FastAPI(title="BI-RADS Mammography Classification API")
-
-@app.on_event("startup")
-async def startup_event():
-    """Wywoływane przy starcie aplikacji FastAPI."""
-    initialize_model()
-
-@app.post("/predict/", response_model=List[Dict[str, Any]])
-async def predict_image(file: UploadFile = File(...)):
-    """
-    Endpoint do klasyfikacji obrazu mammograficznego.
-    Oczekuje pliku obrazu (JPG, PNG).
-    Zwraca listę z wynikami (nawet jeśli tylko jeden obraz).
-    """
-    if model_instance is None or transform_pipeline is None:
-        raise HTTPException(status_code=503, detail="Model nie jest zainicjalizowany. Spróbuj ponownie za chwilę.")
-
-    try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Nie można odczytać pliku obrazu: {e}")
-
-    # Preprocessing
-    input_tensor = transform_pipeline(image).unsqueeze(0).to(DEVICE)
-
-    # Inference
-    with torch.no_grad():
-        model_outputs = model_instance(input_tensor)
-    
-    # Postprocessing
-    probs = torch.nn.functional.softmax(model_outputs, dim=1)
-    confidences, predicted_indices = torch.max(probs, 1)
-    
-    results = []
-    for i in range(len(predicted_indices)): # Pętla na wypadek przyszłego batch processingu
-        birads_category = predicted_indices[i].item() + 1
-        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
-        try:
-            for param in model_instance.parameters():
-                param.requires_grad_(True)
-            model_instance.eval()
-
-            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, :]
-                grad_cam_map_serialized = grad_cam_map_np.tolist()
-
-        except Exception as e:
-            print(f"Błąd podczas generowania Grad-CAM w API: {e}")
-
-        results.append({
-            "birads": birads_category,
-            "confidence": confidence,
-            "interpretation": interpretation,
-            "class_probabilities": class_probabilities,
-            "grad_cam_map": grad_cam_map_serialized
-        })
-    
-    return JSONResponse(content=results)
-
-@app.get("/")
-async def root():
-    return {"message": "Witaj w BI-RADS Classification API! Użyj endpointu /predict/ do wysyłania obrazów."}
-
-# Do uruchomienia lokalnie: uvicorn api_app:app --reload
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+import torch
+import torchvision.models as models
+import torchvision.transforms as transforms
+from PIL import Image
+import torch.nn as nn
+import io
+import numpy as np
+import os
+from typing import List, Dict, Any
+
+# Importy dla Grad-CAM
+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
+
+# --- Konfiguracja ---
+# Upewnij się, że te wartości są zgodne z Twoim repozytorium modelu
+HF_MODEL_REPO_ID = "Enterwar99/MODEL_MAMMOGRAFII"
+MODEL_FILENAME = "best_model.pth" # Nazwa pliku modelu w repozytorium
+
+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]
+
+# Globalne zmienne dla modelu i transformacji
+model_instance = None
+transform_pipeline = None
+
+interpretations_dict = {
+    1: "Wynik negatywny - brak zmian nowotworowych",
+    2: "Zmiana łagodna",
+    3: "Prawdopodobnie zmiana łagodna - zalecana kontrola",
+    4: "Podejrzenie zmiany złośliwej - zalecana biopsja",
+    5: "Wysoka podejrzliwość złośliwości - wymagana biopsja"
+}
+
+# --- Inicjalizacja modelu ---
+def initialize_model():
+    global model_instance, transform_pipeline
+    if model_instance is not None:
+        return
+
+    print(f"Pobieranie modelu {MODEL_FILENAME} z repozytorium {HF_MODEL_REPO_ID}...")
+    try:
+        # Odczytaj token z sekretów, jeśli jest dostępny
+        # Nazwa zmiennej środowiskowej musi być taka sama jak nazwa sekretu w ustawieniach Space
+        hf_auth_token = os.environ.get("HF_TOKEN_MODEL_READ")
+        if hf_auth_token:
+            print("Używam tokenu HF_TOKEN_MODEL_READ do pobrania modelu.")
+        else:
+            print("OSTRZEŻENIE: Sekret HF_TOKEN_MODEL_READ nie został znaleziony. Próba pobrania modelu bez tokenu (może się nie udać dla prywatnych repozytoriów).")
+
+        model_pt_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=MODEL_FILENAME, token=hf_auth_token)
+    except Exception as e:
+        print(f"Błąd podczas pobierania modelu z Hugging Face Hub: {e}")
+        raise RuntimeError(f"Nie można pobrać modelu: {e}")
+
+    print(f"Inicjalizacja architektury modelu ResNet-18...")
+    model_arch = models.resnet18(weights=None)
+    num_feats = model_arch.fc.in_features
+    model_arch.fc = nn.Sequential(
+        nn.Dropout(0.5),
+        nn.Linear(num_feats, 5)
+    )
+    
+    print(f"Ładowanie wag modelu z {model_pt_path}...")
+    model_arch.load_state_dict(torch.load(model_pt_path, map_location=DEVICE))
+    model_arch.to(DEVICE)
+    model_arch.eval()
+    model_instance = model_arch
+
+    transform_pipeline = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+    ])
+    print(f"Model BI-RADS classifier initialized successfully on device: {DEVICE}")
+
+# --- Aplikacja FastAPI ---
+app = FastAPI(title="BI-RADS Mammography Classification API")
+
+@app.on_event("startup")
+async def startup_event():
+    """Wywoływane przy starcie aplikacji FastAPI."""
+    initialize_model()
+
+@app.post("/predict/", response_model=List[Dict[str, Any]])
+async def predict_image(file: UploadFile = File(...)):
+    """
+    Endpoint do klasyfikacji obrazu mammograficznego.
+    Oczekuje pliku obrazu (JPG, PNG).
+    Zwraca listę z wynikami (nawet jeśli tylko jeden obraz).
+    """
+    if model_instance is None or transform_pipeline is None:
+        raise HTTPException(status_code=503, detail="Model nie jest zainicjalizowany. Spróbuj ponownie za chwilę.")
+
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Nie można odczytać pliku obrazu: {e}")
+
+    # Preprocessing
+    input_tensor = transform_pipeline(image).unsqueeze(0).to(DEVICE)
+
+    # Inference
+    with torch.no_grad():
+        model_outputs = model_instance(input_tensor)
+    
+    # Postprocessing
+    probs = torch.nn.functional.softmax(model_outputs, dim=1)
+    confidences, predicted_indices = torch.max(probs, 1)
+    
+    results = []
+    for i in range(len(predicted_indices)): # Pętla na wypadek przyszłego batch processingu
+        birads_category = predicted_indices[i].item() + 1
+        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
+        try:
+            for param in model_instance.parameters():
+                param.requires_grad_(True)
+            model_instance.eval()
+
+            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, :]
+                grad_cam_map_serialized = grad_cam_map_np.tolist()
+
+        except Exception as e:
+            print(f"Błąd podczas generowania Grad-CAM w API: {e}")
+
+        results.append({
+            "birads": birads_category,
+            "confidence": confidence,
+            "interpretation": interpretation,
+            "class_probabilities": class_probabilities,
+            "grad_cam_map": grad_cam_map_serialized
+        })
+    
+    return JSONResponse(content=results)
+
+@app.get("/")
+async def root():
+    return {"message": "Witaj w BI-RADS Classification API! Użyj endpointu /predict/ do wysyłania obrazów."}
+
+# Do uruchomienia lokalnie: uvicorn api_app:app --reload