switched to segformer fashion

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from transformers import SamModel, SamProcessor
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
 import torch
 from PIL import Image
 import numpy as np
@@ -14,11 +14,11 @@ app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Carica il modello e il processore SAM
+# Carica il modello e il processore SegFormer
 try:
-    logger.info("Caricamento del modello SAM...")
-    model = SamModel.from_pretrained("facebook/sam-vit-base")
-    processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+    logger.info("Caricamento del modello SegFormer...")
+    model = SegformerForSemanticSegmentation.from_pretrained("sayeed99/segformer-b3-fashion")
+    processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer-b3-fashion")
     model.to("cpu")  # Usa CPU per il free tier
     logger.info("Modello caricato con successo.")
 except Exception as e:
@@ -27,23 +27,23 @@ except Exception as e:
 
 # Funzione per segmentare l'immagine
 def segment_image(image: Image.Image):
-    # Prepara l'input per SAM
+    # Prepara l'input per SegFormer
     logger.info("Preparazione dell'immagine per l'inferenza...")
-    inputs = processor(image, return_tensors="pt").to("cpu")
+    inputs = processor(images=image, return_tensors="pt").to("cpu")
     
     # Inferenza
     logger.info("Esecuzione dell'inferenza...")
     with torch.no_grad():
-        outputs = model(**inputs, multimask_output=False)
-    
+        outputs = model(**inputs)
+        logits = outputs.logits
+        
     # Post-processa la maschera
     logger.info("Post-processing della maschera...")
-    mask = processor.image_processor.post_process_masks(
-        outputs.pred_masks, inputs["original_sizes"], inputs["reshaped_input_sizes"]
-    )[0][0].cpu().numpy()
+    mask = torch.argmax(logits, dim=1)[0]
+    mask = mask.cpu().numpy()
     
     # Converti la maschera in immagine
-    mask_img = Image.fromarray((mask * 255).astype(np.uint8))
+    mask_img = Image.fromarray((mask * 255 / mask.max()).astype(np.uint8))
     
     # Converti la maschera in base64 per la risposta
     buffered = io.BytesIO()
@@ -51,7 +51,7 @@ def segment_image(image: Image.Image):
     mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
     
     # Annotazioni
-    annotations = {"mask": mask.tolist(), "label": "object"}
+    annotations = {"mask": mask.tolist(), "label": "fashion"}
     
     return mask_base64, annotations
 
@@ -74,7 +74,7 @@ async def segment_endpoint(file: UploadFile = File(...)):
         logger.error(f"Errore nell'endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Errore nell'elaborazione: {str(e)}")
 
-# Per compatibilità con Hugging Face Spaces (Uvicorn viene gestito automaticamente)
+# Per compatibilità con Hugging Face Spaces
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)