Add: SWIN-Transformer-Model-Deployment

Dockerfile

@@ -0,0 +1,16 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--allow-websocket-origin", "youssefabdelmottaleb-Garbage-Classification.hf.space"]

app.py

@@ -0,0 +1,41 @@
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from PIL import Image
+import torch
+import logging
+from models import GarbageClassifier
+
+app = FastAPI()
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load model and processor
+model_dir = "youssefabdelmottaleb/Garbage-Classification-SWIN-Transformer"
+classifier = GarbageClassifier(model_dir)
+
+
+# Endpoint to receive images and return predictions
+@app.post("/predict")
+async def predict_endpoint(file: UploadFile = File(...)):
+    try:
+        image = Image.open(file.file).convert("RGB")
+        inputs = classifier.processor(images=image, return_tensors="pt")
+
+        with torch.no_grad():
+            outputs = classifier.model(**inputs)
+
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+        predicted_class = classifier.labels[predicted_class_idx]
+
+        return {"class": predicted_class}
+
+    except Exception as e:
+        logger.error(f"Error processing image: {e}")
+        raise HTTPException(status_code=500, detail="Error processing image")
+
+
+image_path = "C:/Users/youss/Downloads/paperr.jpg"
+classifier.evaluate_image(image_path)
+
+# To run the app: uvicorn app:app --host 0.0.0.0 --port 8000

models.py

@@ -0,0 +1,31 @@
+import torch
+from transformers import SwinForImageClassification, AutoImageProcessor
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from transformers import SwinForImageClassification, AutoImageProcessor
+
+class GarbageClassifier:
+    def __init__(self, model_dir, num_labels=4):
+        self.labels = ['glass', 'metal', 'paper', 'plastic']
+        self.model, self.processor = self.load_model_and_processor(model_dir, num_labels)
+
+    def load_model_and_processor(self, model_dir, num_labels):
+        model = SwinForImageClassification.from_pretrained(model_dir, num_labels=num_labels)
+        processor = AutoImageProcessor.from_pretrained(model_dir)
+        return model, processor
+
+
+    def evaluate_image(self, image_path):
+        image = Image.open(image_path)
+        inputs = self.processor(images=image, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            predicted_class = logits.argmax(-1).item()
+        
+        plt.imshow(image)
+        plt.title(f"Predicted Class: {self.labels[predicted_class]}")
+        plt.axis('off')
+        plt.show()

predict.py

@@ -0,0 +1,89 @@
+from PIL import Image
+import torch
+import os
+import threading
+import time
+from dotenv import load_dotenv
+from azure.storage.blob import BlobServiceClient
+from transformers import SwinForImageClassification, AutoImageProcessor
+from models import GarbageClassifier
+
+CURRENT_DIR = os.getcwd()
+ROOT_DIR = os.path.abspath(os.path.join(CURRENT_DIR, ".."))
+
+load_dotenv()
+
+# Define Azure Blob Storage settings and model path
+AZURE_CONNECTION_STRING = os.environ.get("AZURE_CONNECTION_STRING")
+AZURE_CONTAINER_NAME = os.environ.get("AZURE_CONTAINER_NAME")
+MODEL_PATH = os.environ.get("MODEL_PATH", "youssefabdelmottaleb/Garbage-Classification-SWIN-Transformer")
+
+# Initialize the Swin Transformer model and processor
+model = SwinForImageClassification.from_pretrained(MODEL_PATH)
+processor = AutoImageProcessor.from_pretrained(MODEL_PATH)
+
+# Function to save image to Azure Blob Storage
+def save_image_to_azure(image, result_class):
+    # Convert PIL image to bytes
+    from io import BytesIO
+    image_bytes = BytesIO()
+    image.save(image_bytes, format='JPEG')
+    image_bytes = image_bytes.getvalue()
+
+    # Create the BlobServiceClient object
+    blob_service_client = BlobServiceClient.from_connection_string(AZURE_CONNECTION_STRING)
+
+    # Create a unique blob name
+    blob_name = f"predicted_images/{result_class}/{result_class}_{int(time.time())}.jpg"
+
+    # Create a blob client using the local file name as the name for the blob
+    blob_client = blob_service_client.get_blob_client(container=AZURE_CONTAINER_NAME, blob=blob_name)
+
+    # Upload the created file
+    blob_client.upload_blob(image_bytes)
+    print(f"Image saved to Azure Blob Storage: {blob_name}")
+
+# Function to save image to local path
+def save_image_to_local_path(image, result_class):
+    # Define the directory based on the class
+    directory = os.path.join('./predicted_images', result_class)
+    os.makedirs(directory, exist_ok=True)
+    
+    # Define the image path
+    image_path = os.path.join(directory, f"{result_class}_{int(time.time())}.jpg")
+    
+    # Save the image
+    image.save(image_path)
+    print(f"Image saved to {image_path}")
+
+# Function to predict using Swin Transformer model
+def predict(image_file):
+    image = Image.open(image_file).convert('RGB')  # Ensure image is RGB
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Preprocess the image
+    inputs = processor(images=image, return_tensors="pt")
+    inputs = {k: v.unsqueeze(0).to(device) for k, v in inputs.items()}
+    
+    # Perform inference
+    model.to(device)
+    model.eval()
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Post-process the outputs to get the predicted class
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    predicted_class = processor.labels[predicted_class_idx]
+
+    # Start a thread to save the image based on the result
+    thread = threading.Thread(target=save_image_to_azure, args=(image, predicted_class))
+    thread.start()
+
+    return {"class": predicted_class}
+
+if __name__ == "__main__":
+    image_file = os.path.join(ROOT_DIR, "input/Ecomate_Dataset/metal/ecomate_metal_41.jpg")
+    result = predict(image_file)
+    print("Prediction result:", result)

requirements.txt

@@ -0,0 +1,12 @@
+fastapi
+uvicorn[standard]
+
+numpy==1.26.3
+torch==2.2.0
+torchvision
+fastapi
+uvicorn
+pillow
+python-multipart
+azure-storage-blob
+uvicorn==0.30.1