updated to onnx

.gitignore

@@ -0,0 +1,5 @@
+.env 
+.DS_Store 
+models/
+model_cache
+onnx_models

README.md

@@ -51,4 +51,29 @@ curl -X POST "https://alexgenovese-segmentation.hf.space/segment-url" \
 -d '{
     "url": "https://plus.unsplash.com/premium_photo-1673210886161-bfcc40f54d1f?ixlib=rb-4.0.3&ixid=MnwxMjA3fDB8MHxzZWFyY2h8MXx8cGVyc29uJTIwc3RhbmRpbmd8ZW58MHx8MHx8&w=1000&q=80"
 }' \
--o response.json
+-o response.json
+
+
+# Segment-clothes-url 
+
+curl -X POST "https://alexgenovese-segmentation.hf.space/segment-clothes-url" \
+-H "Content-Type: application/json" \
+-d '{
+    "url": "https://plus.unsplash.com/premium_photo-1673210886161-bfcc40f54d1f?ixlib=rb-4.0.3&ixid=MnwxMjA3fDB8MHxzZWFyY2h8MXx8cGVyc29uJTIwc3RhbmRpbmd8ZW58MHx8MHx8&w=1000&q=80"
+}' \
+-o response.json
+
+# Convert to ONNX file 
+
+# For the fashion segmentation model:
+python convert_to_onnx.py --model "sayeed99/segformer-b3-fashion" --output "models/fashion_segformer.onnx"
+
+# For the clothes segmentation model:
+python convert_to_onnx.py --model "mattmdjaga/segformer_b2_clothes" --output "models/clothes_segformer.onnx"
+
+
+# Convert To Onnx file 
+
+python3 hf_onnx_converter.py \
+  --source "mattmdjaga/segformer_b2_clothes" \
+  --target "alexgenovese/segformer-onnx"

app.py

@@ -1,193 +1,148 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor, AutoModelForSemanticSegmentation
+from transformers import SegformerImageProcessor
+from huggingface_hub import hf_hub_download
 from pydantic import BaseModel
 from PIL import Image
 import numpy as np
-import io, base64, logging, requests, torch 
-import torch.nn as nn
+import io, base64, logging, requests, os
+import onnxruntime as ort
+from dotenv import load_dotenv
 
-# Inizializza l'app FastAPI
-app = FastAPI()
+# Load environment variables
+load_dotenv()
 
-# Add this class for the request body
 class ImageURL(BaseModel):
     url: str
 
-# Configura il logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Carica il modello e il processore SegFormer
-try:
-    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:
-    logger.error(f"Errore nel caricamento del modello: {str(e)}")
-    raise RuntimeError(f"Errore nel caricamento del modello: {str(e)}")
-
+class ModelManager:
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+        self.token = os.getenv("HF_TOKEN")
+        if not self.token:
+            raise ValueError("HF_TOKEN environment variable is required")
+        self._initialize_models()
+
+    def _initialize_models(self):
+        try:
+            # Initialize ONNX runtime sessions
+            self.logger.info("Loading ONNX models...")
+            
+            # Download and load fashion model
+            fashion_path = hf_hub_download(
+                repo_id="alexgenovese/segformer-onnx",
+                filename="segformer-b3-fashion.onnx",
+                token=self.token
+            )
+            self.fashion_model = ort.InferenceSession(fashion_path)
+            self.fashion_processor = SegformerImageProcessor.from_pretrained(
+                "sayeed99/segformer-b3-fashion",
+                token=self.token
+            )
+
+            # Download and load clothes model
+            clothes_path = hf_hub_download(
+                repo_id="alexgenovese/segformer-onnx",
+                filename="segformer_b2_clothes.onnx",
+                token=self.token
+            )
+            self.clothes_model = ort.InferenceSession(clothes_path)
+            self.clothes_processor = SegformerImageProcessor.from_pretrained(
+                "mattmdjaga/segformer_b2_clothes",
+                token=self.token
+            )
+
+            self.logger.info("All models loaded successfully.")
+        except Exception as e:
+            self.logger.error(f"Error initializing models: {str(e)}")
+            raise RuntimeError(f"Error initializing models: {str(e)}")
+
+    def process_fashion_image(self, image: Image.Image):
+        inputs = self.fashion_processor(images=image, return_tensors="np")
+        onnx_inputs = {
+            'input': inputs['pixel_values']
+        }
+        logits = self.fashion_model.run(None, onnx_inputs)[0]
+        return self._post_process_outputs(logits, image.size)
 
-# Add new model and processor initialization after existing ones
-try:
-    logger.info("Loading clothes segmentation model...")
-    clothes_model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
-    clothes_processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
-    clothes_model.to("cpu")
-    logger.info("Clothes model loaded successfully.")
-except Exception as e:
-    logger.error(f"Error loading clothes model: {str(e)}")
-    raise RuntimeError(f"Error loading clothes model: {str(e)}")
+    def process_clothes_image(self, image: Image.Image):
+        inputs = self.clothes_processor(images=image, return_tensors="np")
+        onnx_inputs = {
+            'input': inputs['pixel_values']
+        }
+        logits = self.clothes_model.run(None, onnx_inputs)[0]
+        return self._post_process_outputs(logits, image.size)
 
-# Funzione per segmentare l'immagine
-def segment_image(image: Image.Image):
-    # Prepara l'input per SegFormer
-    logger.info("Preparazione dell'immagine per l'inferenza...")
-    inputs = processor(images=image, return_tensors="pt").to("cpu")
-    
-    # Inferenza
-    logger.info("Esecuzione dell'inferenza...")
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+    def _post_process_outputs(self, logits, image_size):
+        # Convert logits to proper shape for processing
+        logits = np.array(logits)
         
-    # Post-processa la maschera
-    logger.info("Post-processing della maschera...")
-    mask = torch.argmax(logits, dim=1)[0]
-    mask = mask.cpu().numpy()
-    
-    # Converti la maschera in immagine
-    mask_img = Image.fromarray((mask * 255 / mask.max()).astype(np.uint8))
-    
-    # Converti la maschera in base64 per la risposta
-    buffered = io.BytesIO()
-    mask_img.save(buffered, format="PNG")
-    mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    
-    # Annotazioni
-    annotations = {"mask": mask.tolist(), "label": logits }
-    
-    return mask_base64, annotations
-
-# Endpoint API
-@app.post("/segment")
-async def segment_endpoint(file: UploadFile = File(...)):
-    try:
-        logger.info("Ricezione del file...")
-        image_data = await file.read()
-        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        # Resize prediction to match original image size
+        from skimage.transform import resize
+        resized_logits = resize(
+            logits[0],
+            (image_size[1], image_size[0]),
+            order=1,
+            preserve_range=True,
+            mode='reflect'
+        )
         
-        logger.info("Segmentazione dell'immagine...")
-        mask_base64, annotations = segment_image(image)
+        # Get prediction
+        pred_seg = np.argmax(resized_logits, axis=0)
+        mask_img = Image.fromarray((pred_seg * 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}",
-            "annotations": annotations
+            "size": image_size,
+            "predictions": pred_seg.tolist()
         }
-    except Exception as e:
-        logger.error(f"Errore nell'endpoint: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Errore nell'elaborazione: {str(e)}")
-
 
+# Initialize FastAPI and ModelManager
+app = FastAPI()
+model_manager = ModelManager()
 
-# 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
-        }
+        return model_manager.process_fashion_image(image)
         
     except Exception as e:
-        logger.error(f"Error processing URL: {str(e)}")
+        logging.error(f"Error processing URL: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 
-
-
-# Add new endpoint
 @app.post("/segment-clothes-url")
 async def segment_clothes_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 = clothes_processor(images=image, return_tensors="pt")
-        outputs = clothes_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,
-            "predictions": pred_seg.numpy().tolist()
-        }
+        return model_manager.process_clothes_image(image)
         
     except Exception as e:
-        logger.error(f"Error processing URL: {str(e)}")
+        logging.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
+@app.post("/segment")
+async def segment_endpoint(file: UploadFile = File(...)):
+    try:
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        return model_manager.process_fashion_image(image)
+    except Exception as e:
+        logging.error(f"Error in endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing: {str(e)}")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)

hf_onnx_converter.py

@@ -0,0 +1,202 @@
+import torch
+from transformers import AutoModelForSemanticSegmentation, SegformerImageProcessor
+from huggingface_hub import HfApi, create_repo, upload_file, model_info
+import os
+from dotenv import load_dotenv
+from pathlib import Path
+import logging
+import argparse
+import tempfile
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv()
+
+class ConfigurationError(Exception):
+    """Raised when required environment variables are missing"""
+    pass
+
+class HFOnnxConverter:
+    def __init__(self, token=None):
+        # Load configuration from environment
+        self.token = token or os.getenv("HF_TOKEN")
+        self.model_cache_dir = os.getenv("MODEL_CACHE_DIR")
+        self.onnx_output_dir = os.getenv("ONNX_OUTPUT_DIR")
+
+        # Validate configuration
+        if not self.token:
+            raise ConfigurationError("HF_TOKEN is required in environment variables")
+
+        # Create directories if they don't exist
+        for directory in [self.model_cache_dir, self.onnx_output_dir]:
+            if directory:
+                Path(directory).mkdir(parents=True, exist_ok=True)
+
+        self.api = HfApi()
+
+        # Login to Hugging Face
+        try:
+            self.api.whoami(token=self.token)
+            logger.info("Successfully authenticated with Hugging Face")
+        except Exception as e:
+            raise ConfigurationError(f"Failed to authenticate with Hugging Face: {str(e)}")
+        
+    def setup_repository(self, repo_name: str) -> str:
+        """Create or get repository on Hugging Face Hub"""
+        try:
+            create_repo(
+                repo_name,
+                token=self.token,
+                private=False,
+                exist_ok=True
+            )
+            logger.info(f"Repository {repo_name} is ready")
+            return repo_name
+        except Exception as e:
+            logger.error(f"Error setting up repository: {e}")
+            raise
+
+    def verify_model_exists(self, model_name: str) -> bool:
+        """Verify if the model exists and is accessible"""
+        try:
+            model_info(model_name, token=self.token)
+            return True
+        except Exception as e:
+            logger.error(f"Model verification failed: {str(e)}")
+            return False
+
+    def convert_and_push(self, source_model: str, target_repo: str):
+        """Convert model to ONNX and push to Hugging Face Hub"""
+        try:
+            # Verify model exists and is accessible
+            if not self.verify_model_exists(source_model):
+                raise ValueError(f"Model {source_model} is not accessible. Check if the model exists and you have proper permissions.")
+
+            # Use model cache directory if specified
+            model_kwargs = {
+                "token": self.token
+            }
+            if self.model_cache_dir:
+                model_kwargs["cache_dir"] = self.model_cache_dir
+
+            # Create working directory
+            working_dir = self.onnx_output_dir or tempfile.mkdtemp()
+            tmp_path = Path(working_dir) / f"{target_repo.split('/')[-1]}.onnx"
+                
+            logger.info(f"Loading model {source_model}...")
+            model = AutoModelForSemanticSegmentation.from_pretrained(
+                source_model,
+                **model_kwargs
+            )
+            processor = SegformerImageProcessor.from_pretrained(
+                source_model,
+                **model_kwargs
+            )
+
+            # Set model to evaluation mode
+            model.eval()
+
+            # Create dummy input
+            dummy_input = processor(
+                images=torch.zeros(1, 3, 224, 224),
+                return_tensors="pt"
+            )
+
+            # Export to ONNX
+            logger.info(f"Converting to ONNX format... Output path: {tmp_path}")
+            torch.onnx.export(
+                model,
+                (dummy_input['pixel_values'],),
+                tmp_path,
+                input_names=['input'],
+                output_names=['output'],
+                dynamic_axes={
+                    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
+                    'output': {0: 'batch_size'}
+                },
+                opset_version=12,
+                do_constant_folding=True
+            )
+
+            # Create model card with environment info
+            model_card = f"""---
+base_model: {source_model}
+tags:
+- onnx
+- semantic-segmentation
+---
+
+# ONNX Model converted from {source_model}
+
+This is an ONNX version of the model {source_model}, converted automatically.
+
+## Model Information
+- Original Model: {source_model}
+- ONNX Opset Version: 12
+- Input Shape: Dynamic (batch_size, 3, height, width)
+
+## Usage
+
+```python
+import onnxruntime as ort
+import numpy as np
+
+# Load ONNX model
+session = ort.InferenceSession("model.onnx")
+
+# Prepare input
+input_data = np.zeros((1, 3, 224, 224), dtype=np.float32)
+
+# Run inference
+outputs = session.run(None, {{"input": input_data}})
+```
+"""
+            # Save model card
+            readme_path = Path(working_dir) / "README.md"
+            with open(readme_path, "w") as f:
+                f.write(model_card)
+
+            # Push files to hub
+            logger.info(f"Pushing files to {target_repo}...")
+            self.api.upload_file(
+                path_or_fileobj=str(tmp_path),
+                path_in_repo="model.onnx",
+                repo_id=target_repo,
+                token=self.token
+            )
+            self.api.upload_file(
+                path_or_fileobj=str(readme_path),
+                path_in_repo="README.md",
+                repo_id=target_repo,
+                token=self.token
+            )
+
+            logger.info(f"Successfully pushed ONNX model to {target_repo}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error during conversion and upload: {e}")
+            return False
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert and push model to ONNX format on Hugging Face Hub')
+    parser.add_argument('--source', type=str, required=True,
+                      help='Source model name (e.g., "sayeed99/segformer-b3-fashion")')
+    parser.add_argument('--target', type=str, required=True,
+                      help='Target repository name (e.g., "your-username/model-name-onnx")')
+    parser.add_argument('--token', type=str, help='Hugging Face token (optional)')
+
+    args = parser.parse_args()
+
+    converter = HFOnnxConverter(token=args.token)
+    converter.setup_repository(args.target)
+    success = converter.convert_and_push(args.source, args.target)
+    
+    if not success:
+        exit(1)
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -4,4 +4,7 @@ torch
 torchvision
 transformers
 pillow
-numpy
+numpy
+torch 
+dotenv
+onnx