added new endpoint

README.md

@@ -37,4 +37,8 @@ ir
 
 curl -X POST "http://localhost:7860/segment" \
 -H "Content-Type: application/json" \
--d "{\"image_base64\": \"$(base64 woman_with_bag.jpeg)\"}"
+-d "{\"image_base64\": \"$(base64 woman_with_bag.jpeg)\"}"
+
+
+# Output 
+{"mask":"data:image/png;base64...","annotations":{"mask":[[]]"label":"fashion"}} 

app.py

@@ -1,15 +1,22 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
 import torch
+from pydantic import BaseModel
 from PIL import Image
 import numpy as np
 import io
 import base64
 import logging
+import requests
+import torch.nn as nn
 
 # Inizializza l'app FastAPI
 app = FastAPI()
 
+# Add this class for the request body
+class ImageURL(BaseModel):
+    url: str
+
 # Configura il logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -74,6 +81,56 @@ 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)}")
 
+
+
+# Add new endpoint
+@app.post("/segment-url")
+async def segment_url_endpoint(image_data: ImageURL):
+    try:
+        logger.info("Downloading image from URL...")
+        response = requests.get(image_data.url, stream=True)
+        if response.status_code != 200:
+            raise HTTPException(status_code=400, detail="Could not download image from URL")
+        
+        # Open image from URL
+        image = Image.open(response.raw).convert("RGB")
+        
+        # Process image with SegFormer
+        logger.info("Processing image...")
+        inputs = processor(images=image, return_tensors="pt")
+        outputs = model(**inputs)
+        logits = outputs.logits.cpu()
+
+        # Upsample logits to match original image size
+        upsampled_logits = nn.functional.interpolate(
+            logits,
+            size=image.size[::-1],
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        # Get prediction
+        pred_seg = upsampled_logits.argmax(dim=1)[0]
+        
+        # Convert to image
+        mask_img = Image.fromarray((pred_seg.numpy() * 255).astype(np.uint8))
+        
+        # Convert to base64
+        buffered = io.BytesIO()
+        mask_img.save(buffered, format="PNG")
+        mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        return {
+            "mask": f"data:image/png;base64,{mask_base64}",
+            "size": image.size,
+            "labels" : pred_seg
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing URL: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
+
 # Per compatibilità con Hugging Face Spaces
 if __name__ == "__main__":
     import uvicorn