Tutilisation d'un fichier de fonction allégé pour l'API

app.py

@@ -2,7 +2,8 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import StreamingResponse
 import torch
-from fonctions import charger_segformer
+from fonctions_api import charger_segformer, remap_classes
+from fonctions_api import decode_cityscapes_mask
 from PIL import Image
 import io
 import numpy as np

fonctions_api.py

@@ -0,0 +1,58 @@
+# fonctions_api.py (version allégée pour l'API FastAPI avec SegFormer)
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from pathlib import Path
+from transformers import SegformerForSemanticSegmentation
+
+# -------- Palette (optionnelle pour colorisation) --------
+PALETTE = {
+    0: (0, 0, 0),           # void
+    1: (50, 50, 150),       # flat
+    2: (102, 0, 204),       # construction
+    3: (255, 85, 0),        # object
+    4: (255, 255, 0),       # nature
+    5: (0, 255, 255),       # sky
+    6: (255, 0, 255),       # human
+    7: (255, 255, 255),     # vehicle
+}
+
+# -------- Fonction principale pour charger SegFormer --------
+def charger_segformer(num_classes=8):
+    model = SegformerForSemanticSegmentation.from_pretrained(
+        "nvidia/segformer-b5-finetuned-ade-640-640",
+        num_labels=num_classes,
+        ignore_mismatched_sizes=True
+    )
+    model.config.num_labels = num_classes
+    model.config.output_hidden_states = False
+    return model
+
+# -------- Remapping Cityscapes labelIds vers 8 classes --------
+def remap_classes(mask: np.ndarray) -> np.ndarray:
+    labelIds_to_main_classes = {
+        0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0, 6: 0,
+        7: 1, 8: 1,
+        9: 0, 10: 0,
+        11: 2, 12: 2, 13: 2,
+        14: 0, 15: 0, 16: 0,
+        17: 3, 18: 3, 19: 3, 20: 3,
+        21: 4, 22: 4,
+        23: 5,
+        24: 6, 25: 6,
+        26: 7, 27: 7, 28: 7, 29: 7, 30: 7, 31: 7, 32: 7, 33: 7
+    }
+    remapped_mask = np.copy(mask)
+    for original_class, new_class in labelIds_to_main_classes.items():
+        remapped_mask[mask == original_class] = new_class
+    remapped_mask[mask > 33] = 0
+    return remapped_mask.astype(np.uint8)
+
+# -------- Convertit un masque 2D en image RGB (optionnel) --------
+def decode_cityscapes_mask(mask):
+    h, w = mask.shape
+    mask_rgb = np.zeros((h, w, 3), dtype=np.uint8)
+    for class_id, color in PALETTE.items():
+        mask_rgb[mask == class_id] = color
+    return mask_rgb