Add Application Files

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.11.0
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+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"]

app.py

@@ -0,0 +1,105 @@
+from fastapi import FastAPI, Request, UploadFile, File, HTTPException
+import uvicorn
+from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.templating import Jinja2Templates
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from os import path, makedirs
+from PIL import Image
+import logging
+from io import BytesIO
+from typing import Tuple
+from utils.image_segmenter import ImageSegmenter
+import zipfile
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="Image Background Remover",
+             description="API for removing image backgrounds using ML",
+             version="1.0.0")
+
+# Define allowed origins explicitly for security
+origins = [
+    "http://127.0.0.1:5500",
+    "http://localhost:5500"
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,  # Use the defined origins instead of "*"
+    allow_credentials=True,
+    allow_methods=["GET", "POST"],  # Specify only needed methods
+    allow_headers=["*"],
+)
+
+# Set up the templates
+templates = Jinja2Templates(directory="templates")
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# Initialize ImageSegmenter once at startup
+segmenter = ImageSegmenter()
+
+@app.get("/")
+async def index(request: Request) -> HTMLResponse:
+    return templates.TemplateResponse("main.html", {"request": request})
+
+@app.post("/api/remove-background/")
+async def remove_background(file_obj: UploadFile = File(...)) -> StreamingResponse:
+    try:
+        # Validate file type
+        if not file_obj.content_type.startswith('image/'):
+            raise HTTPException(status_code=400, detail="File must be an image")
+        
+        # Read image with error handling
+        try:
+            image_content = await file_obj.read()
+            image = Image.open(BytesIO(image_content))
+        except Exception as e:
+            logger.error(f"Error reading image: {str(e)}")
+            raise HTTPException(status_code=400, detail="Invalid image file")
+
+        # Process image
+        try:
+            image, mask = await segmenter.segment(image)  # Fixed typo in method name
+        except Exception as e:
+            logger.error(f"Error processing image: {str(e)}")
+            raise HTTPException(status_code=500, detail="Error processing image")
+
+        # Create ZIP file in memory
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+            # Save processed image
+            image_buffer = BytesIO()
+            image.save(image_buffer, "PNG", optimize=True)
+            image_buffer.seek(0)
+            zip_file.writestr('processed_image.png', image_buffer.getvalue())
+            
+            # Save mask
+            mask_buffer = BytesIO()
+            mask.save(mask_buffer, "PNG", optimize=True)
+            mask_buffer.seek(0)
+            zip_file.writestr('mask.png', mask_buffer.getvalue())
+
+        zip_buffer.seek(0)
+        
+        return StreamingResponse(
+            zip_buffer,
+            media_type="application/zip",
+            headers={
+                "Content-Disposition": f"attachment; filename=result_{file_obj.filename}.zip"
+            }
+        )
+    except Exception as e:
+        logger.error(f"Unexpected error: {str(e)}")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+# if __name__ == "__main__":
+#     uvicorn.run(
+#         app, 
+#         host="127.0.0.1", 
+#         port=8000,
+#         log_level="info",
+#         reload=True  # Enable auto-reload during development
+#     )

requirements.txt

@@ -0,0 +1,11 @@
+fastapi
+uvicorn
+pillow
+torch
+torchvision
+transformers
+python-multipart
+jinja2
+aiofiles
+kornia
+pyngrok

static/script.js

@@ -0,0 +1,165 @@
+const uploadArea = document.getElementById("uploadArea");
+const fileInput = document.getElementById("fileInput");
+const previewImage = document.getElementById("previewImage");
+const resultImage = document.getElementById("resultImage");
+const removeBackgroundBtn = document.getElementById("removeBackgroundBtn");
+const uploadLoading = document.getElementById("uploadLoading");
+const resultLoading = document.getElementById("resultLoading");
+
+// Add these new elements to your HTML
+const imageContainer = document.querySelector(".image-box:nth-child(2)");
+const toggleContainer = document.createElement("div");
+toggleContainer.className = "toggle-container";
+toggleContainer.innerHTML = `
+  <select id="viewSelector" class="view-selector">
+    <option value="processed">Processed Image</option>
+    <option value="mask">Mask</option>
+  </select>
+`;
+imageContainer.insertBefore(
+  toggleContainer,
+  imageContainer.querySelector(".upload-area")
+);
+
+// Store the extracted images
+let processedImageUrl = null;
+let maskImageUrl = null;
+
+uploadArea.addEventListener("click", () => {
+  fileInput.click();
+});
+
+uploadArea.addEventListener("dragover", (e) => {
+  e.preventDefault();
+  uploadArea.style.borderColor = "#666666";
+  uploadArea.style.backgroundColor = "#333333";
+});
+
+uploadArea.addEventListener("dragleave", (e) => {
+  e.preventDefault();
+  uploadArea.style.borderColor = "#444444";
+  uploadArea.style.backgroundColor = "#222222";
+});
+
+uploadArea.addEventListener("drop", (e) => {
+  e.preventDefault();
+  uploadArea.style.borderColor = "#444444";
+  uploadArea.style.backgroundColor = "#222222";
+
+  const file = e.dataTransfer.files[0];
+  if (file && file.type.startsWith("image/")) {
+    handleImageUpload(file);
+  }
+});
+
+fileInput.addEventListener("change", (e) => {
+  const file = e.target.files[0];
+  if (file) {
+    handleImageUpload(file);
+  }
+});
+
+function handleImageUpload(file) {
+  const reader = new FileReader();
+
+  uploadLoading.style.display = "block";
+
+  reader.onload = (e) => {
+    previewImage.src = e.target.result;
+    previewImage.style.display = "block";
+    uploadArea.querySelector("p").style.display = "none";
+    removeBackgroundBtn.disabled = false;
+    uploadLoading.style.display = "none";
+  };
+
+  reader.readAsDataURL(file);
+}
+
+// Function to extract images from ZIP
+async function extractImagesFromZip(blob) {
+  const zip = await JSZip.loadAsync(blob);
+  const files = Object.values(zip.files);
+
+  const processedImageFile = files.find(
+    (file) => file.name === "processed_image.png"
+  );
+  const maskFile = files.find((file) => file.name === "mask.png");
+
+  if (!processedImageFile || !maskFile) {
+    throw new Error("Missing required images in ZIP file");
+  }
+
+  // Get blobs for both images
+  const processedImageBlob = await processedImageFile.async("blob");
+  const maskBlob = await maskFile.async("blob");
+
+  // Create URLs for the images
+  if (processedImageUrl) URL.revokeObjectURL(processedImageUrl);
+  if (maskImageUrl) URL.revokeObjectURL(maskImageUrl);
+
+  processedImageUrl = URL.createObjectURL(processedImageBlob);
+  maskImageUrl = URL.createObjectURL(maskBlob);
+
+  return { processedImageUrl, maskImageUrl };
+}
+
+// Add event listener for view selector
+document.getElementById("viewSelector").addEventListener("change", (e) => {
+  resultImage.src =
+    e.target.value === "processed" ? processedImageUrl : maskImageUrl;
+});
+
+removeBackgroundBtn.addEventListener("click", async () => {
+  try {
+    resultLoading.style.display = "block";
+    removeBackgroundBtn.disabled = true;
+    document.getElementById("viewSelector").disabled = true;
+
+    // Convert base64 to blob
+    const image = fileInput.files[0];
+
+    // Create FormData
+    const formData = new FormData();
+    formData.append("file_obj", image);
+
+    // Send to backend
+    const response = await fetch("/api/remove-background/", {
+      method: "POST",
+      body: formData,
+    });
+
+    if (!response.ok) {
+      throw new Error("Failed to remove background");
+    }
+
+    const blob = await response.blob();
+
+    // Extract and display images from ZIP
+    const { processedImageUrl: processedUrl, maskImageUrl: maskUrl } =
+      await extractImagesFromZip(blob);
+
+    // Show processed image by default
+    resultImage.src = processedUrl;
+    resultImage.style.display = "block";
+    document.getElementById("viewSelector").disabled = false;
+
+    // Create download link
+    const downloadLink = document.createElement("a");
+    downloadLink.href = URL.createObjectURL(blob);
+    downloadLink.download = "images.zip";
+    downloadLink.className = "download-link";
+    downloadLink.innerText = "Download ZIP";
+
+    // Add download link below the image container
+    const downloadContainer = document.createElement("div");
+    downloadContainer.className = "download-container";
+    downloadContainer.appendChild(downloadLink);
+    imageContainer.appendChild(downloadContainer);
+  } catch (error) {
+    console.error("Error:", error);
+    alert("Failed to process image. Please try again.");
+  } finally {
+    resultLoading.style.display = "none";
+    removeBackgroundBtn.disabled = false;
+  }
+});

static/style.css

@@ -0,0 +1,201 @@
+* {
+  margin: 0;
+  padding: 0;
+  box-sizing: border-box;
+  font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
+}
+
+body {
+  min-height: 100vh;
+  background: linear-gradient(135deg, #1a1a1a, #000000);
+  padding: 2rem;
+  color: #ffffff;
+}
+
+.container {
+  max-width: 1200px;
+  margin: 0 auto;
+}
+
+.title {
+  text-align: center;
+  color: #ffffff;
+  margin-bottom: 2rem;
+  font-size: 2rem;
+  text-shadow: 0 2px 4px rgba(0, 0, 0, 0.3);
+}
+
+.image-containers {
+  display: grid;
+  grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
+  gap: 2rem;
+  margin-bottom: 2rem;
+}
+
+.image-box {
+  background-color: #2a2a2a;
+  border-radius: 12px;
+  padding: 1rem;
+  box-shadow: 0 8px 16px rgba(0, 0, 0, 0.4);
+  height: 400px;
+  display: flex;
+  flex-direction: column;
+  position: relative;
+  border: 1px solid #333333;
+}
+
+.image-box h2 {
+  color: #ffffff;
+  margin-bottom: 1rem;
+  font-size: 1.2rem;
+}
+
+.upload-area {
+  flex: 1;
+  border: 2px dashed #444444;
+  border-radius: 8px;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  cursor: pointer;
+  transition: all 0.3s ease;
+  position: relative;
+  overflow: hidden;
+  background-color: #222222;
+}
+
+.upload-area:hover {
+  border-color: #666666;
+  background-color: #333333;
+}
+
+.upload-area p {
+  color: #999999;
+  text-align: center;
+  padding: 1rem;
+}
+
+.preview-image {
+  width: 100%;
+  height: 100%;
+  object-fit: contain;
+  display: none;
+}
+
+.result-image {
+  width: 100%;
+  height: 100%;
+  object-fit: contain;
+}
+
+.remove-bg-btn {
+  display: block;
+  width: 100%;
+  max-width: 300px;
+  margin: 0 auto;
+  padding: 1rem 2rem;
+  background-color: #ffffff;
+  color: #000000;
+  border: none;
+  border-radius: 8px;
+  cursor: pointer;
+  font-size: 1rem;
+  transition: all 0.3s ease;
+  font-weight: 500;
+  box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+}
+
+.remove-bg-btn:hover {
+  background-color: #f0f0f0;
+  transform: translateY(-2px);
+  box-shadow: 0 6px 12px rgba(0, 0, 0, 0.3);
+}
+
+.remove-bg-btn:disabled {
+  background-color: #444444;
+  color: #666666;
+  cursor: not-allowed;
+  transform: none;
+  box-shadow: none;
+}
+
+.loading {
+  position: absolute;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%);
+  display: none;
+}
+
+.loading::after {
+  content: "";
+  width: 40px;
+  height: 40px;
+  border: 4px solid #333333;
+  border-top: 4px solid #ffffff;
+  border-radius: 50%;
+  animation: spin 1s linear infinite;
+  display: block;
+}
+
+@keyframes spin {
+  0% {
+    transform: rotate(0deg);
+  }
+  100% {
+    transform: rotate(360deg);
+  }
+}
+
+@media (max-width: 768px) {
+  body {
+    padding: 1rem;
+  }
+
+  .image-box {
+    height: 300px;
+  }
+
+  .title {
+    font-size: 1.5rem;
+  }
+}
+
+.toggle-container {
+  margin-bottom: 1rem;
+  text-align: center;
+}
+
+.view-selector {
+  padding: 0.5rem;
+  border-radius: 4px;
+  background-color: #333333;
+  color: #ffffff;
+  border: 1px solid #444444;
+  cursor: pointer;
+  width: 150px;
+}
+
+.view-selector:disabled {
+  opacity: 0.5;
+  cursor: not-allowed;
+}
+
+.download-container {
+  text-align: center;
+  margin-top: 1rem;
+}
+
+.download-link {
+  display: inline-block;
+  padding: 0.5rem 1rem;
+  background-color: #333333;
+  color: #ffffff;
+  text-decoration: none;
+  border-radius: 4px;
+  transition: all 0.3s ease;
+}
+
+.download-link:hover {
+  background-color: #444444;
+}

templates/main.html

@@ -0,0 +1,37 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Background Remover</title>
+    <link rel="stylesheet" href="/static/style.css"> 
+
+</head>
+<body>
+    <div class="container">
+        <h1 class="title">Background Remover</h1>
+        <div class="image-containers">
+            <div class="image-box">
+                <h2>Upload Image</h2>
+                <div class="upload-area" id="uploadArea">
+                    <p>Click or drag image here</p>
+                    <input type="file" id="fileInput" hidden accept="image/*">
+                    <img id="previewImage" class="preview-image">
+                    <div class="loading" id="uploadLoading"></div>
+                </div>
+            </div>
+            <div class="image-box">
+                <h2>Result</h2>
+                <div class="upload-area">
+                    <img id="resultImage" class="result-image">
+                    <div class="loading" id="resultLoading"></div>
+                </div>
+            </div>
+        </div>
+        <button class="remove-bg-btn" id="removeBackgroundBtn" disabled>Remove Background</button>
+    </div>
+
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/jszip/3.10.1/jszip.min.js"></script>
+    <script src="/static/script.js"></script>
+</body>
+</html>

utils/image_segmenter.py

@@ -0,0 +1,56 @@
+from PIL import Image
+import torch
+from torchvision import transforms  # type: ignore
+from transformers import AutoModelForImageSegmentation
+from typing import Tuple
+import logging
+
+logger = logging.getLogger(__name__)
+
+class ImageSegmenter:
+    def __init__(self, image_size: Tuple[int, int] = (1024, 1024)):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        logger.info(f"Using device: {self.device}")
+        
+        try:
+            self.birefnet = AutoModelForImageSegmentation.from_pretrained(
+                "ZhengPeng7/BiRefNet", 
+                trust_remote_code=True
+            )
+            self.birefnet.to(self.device)
+            self.birefnet.eval()  # Set model to evaluation mode
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+
+        self.transform_image = transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            ),
+        ])
+    
+    async def extract_object(self, image: Image.Image) -> Tuple[Image.Image, Image.Image]:
+        try:
+            # Transform image
+            input_images = self.transform_image(image).unsqueeze(0).to(self.device)
+
+            # Prediction
+            with torch.no_grad():
+                preds = self.birefnet(input_images)[-1].sigmoid().cpu()
+            
+            pred = preds[0].squeeze()
+            pred_pil = transforms.ToPILImage()(pred)
+            mask = pred_pil.resize(image.size)
+            image.putalpha(mask)
+            
+            return image, mask
+        except Exception as e:
+            logger.error(f"Error in extract_object: {str(e)}")
+            raise
+    
+    async def segment(self, image: Image.Image) -> Tuple[Image.Image, Image.Image]:
+        """Fixed typo in method name and added type hints"""
+        return await self.extract_object(image)