Update api_app.py

api_app.py

@@ -10,6 +10,8 @@ import numpy as np
 import os
 from typing import List, Dict, Any
 import logging # Dodajemy import modułu logging
+import cv2 # Dodajemy OpenCV
+import base64 # Dodajemy base64
 
 # Importy dla Grad-CAM
 from pytorch_grad_cam import GradCAMPlusPlus
@@ -86,6 +88,48 @@ def initialize_model():
     ])
     logger.info(f"Model BI-RADS classifier initialized successfully on device: {DEVICE}")
 
+# --- 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."""
+    logger.info(f"Rozpoczynanie tworzenia obrazu Grad-CAM overlay dla kategorii {birads_category}")
+    try:
+        img_np = np.array(original_pil_image.convert('RGB')).astype(np.float32) / 255.0
+        cam_resized = cv2.resize(grayscale_cam, (img_np.shape[1], img_np.shape[0]))
+
+        cam_normalized = (cam_resized - np.min(cam_resized)) / (np.max(cam_resized) - np.min(cam_resized) + 1e-8)
+
+        threshold = 0.7
+        cam_normalized[cam_normalized < threshold] = 0
+
+        kernel_size = 5
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
+        cam_cleaned = cv2.morphologyEx(cam_normalized, cv2.MORPH_OPEN, kernel)
+
+        birads_colors_rgb = {
+            1: (0.1, 0.7, 0.1),
+            2: (0.53, 0.81, 0.92),
+            3: (1.0, 0.9, 0.0),
+            4: (1.0, 0.5, 0.0),
+            5: (0.9, 0.1, 0.1)
+        }
+        chosen_color = np.array(birads_colors_rgb.get(birads_category, (0.5, 0.5, 0.5)))
+
+        color_overlay_np = np.zeros_like(img_np)
+        for c in range(3):
+            color_overlay_np[:, :, c] = chosen_color[c]
+
+        alpha = cam_cleaned * transparency
+        alpha_expanded = alpha[..., np.newaxis]
+
+        highlighted_image_np = img_np * (1 - alpha_expanded) + color_overlay_np * alpha_expanded
+        highlighted_image_np = np.clip(highlighted_image_np, 0, 1)
+        final_image_np = (highlighted_image_np * 255).astype(np.uint8)
+        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)
+        return None
+
 # --- Aplikacja FastAPI ---
 app = FastAPI(title="BI-RADS Mammography Classification API")
 
@@ -144,6 +188,7 @@ async def predict_image(file: UploadFile = File(...)):
 
         # 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()
@@ -158,8 +203,19 @@ async def predict_image(file: UploadFile = File(...)):
             
             if grayscale_cam is not None:
                 grad_cam_map_np = grayscale_cam[0, :]
-                grad_cam_map_serialized = grad_cam_map_np.tolist()
                 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.")
             else:
                 logger.warning(f"[RequestID: {request_id}] Wygenerowany Grad-CAM jest None.")
         except Exception as e:
@@ -170,7 +226,7 @@ async def predict_image(file: UploadFile = File(...)):
             "confidence": confidence,
             "interpretation": interpretation,
             "class_probabilities": class_probabilities,
-            "grad_cam_map": grad_cam_map_serialized
+            "grad_cam_image_base64": grad_cam_image_base64 # Dodajemy obraz base64
         })
     
     logger.info(f"[RequestID: {request_id}] Przetwarzanie żądania /predict/ zakończone. Zwracam wyniki.")