Initial commit

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+RUN pip install huggingface-hub
+
+RUN huggingface-hub download \
+    apple/MobileCLIP2-S0 \
+    --local-dir /app/model \
+    --local-dir-use-symlinks False \
+    --include "*.pt"
+
+COPY requirements.txt .
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY ./app /app/app
+
+EXPOSE 8000
+
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

app/api/endpoints/staging.py

@@ -0,0 +1,35 @@
+from fastapi import APIRouter, File, UploadFile, Depends, HTTPException
+from PIL import Image
+import io
+from app.models.clip_model import staging_ranker
+from app.schemas.staging import StagingRequest, StagingResponse
+
+router = APIRouter()
+
+@router.post("/rank_image", response_model=StagingResponse)
+async def rank_image_for_staging(
+    prompts: StagingRequest = Depends(),
+    file: UploadFile = File(...)
+):
+    """
+    Accepts an image upload and optional JSON prompts to compute a stageability score.
+    
+    - **file**: The image file to be analyzed.
+    - **prompts**: A JSON object with `prompt_good`, `prompt_bad`, and optional `prompt_aesthetic`.
+    """
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="File provided is not an image.")
+
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+    except Exception:
+        raise HTTPException(status_code=500, detail="Could not process the uploaded image.")
+
+    score = staging_ranker.compute_score(image, prompts)
+    
+    return StagingResponse(
+        filename=file.filename,
+        stageability_score=score,
+        details="Score calculated based on the provided prompts."
+    )

app/core/config.py

@@ -0,0 +1,14 @@
+from pydantic import BaseSettings
+
+class Settings(BaseSettings):
+    """
+    Configuration settings for the application.
+    The model path is defined here, pointing to where the Dockerfile will download it.
+    """
+    MODEL_PATH: str = "/app/model/mobileclip2_s0.pt"
+    MODEL_NAME: str = "MobileCLIP2-S0"
+
+    class Config:
+        env_file = ".env"
+
+settings = Settings()

app/main.py

@@ -0,0 +1,14 @@
+from fastapi import FastAPI
+from app.api.endpoints import staging
+
+app = FastAPI(
+    title="Virtual Staging Ranker API",
+    description="An API to rank images based on their suitability for virtual staging using MobileCLIP2.",
+    version="1.0.0"
+)
+
+app.include_router(staging.router, prefix="/api/v1", tags=["Staging"])
+
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to the Staging Ranker API. Visit /docs for more info."}

app/models/clip_model.py

@@ -0,0 +1,60 @@
+import torch
+import open_clip
+from PIL import Image
+from mobileclip.modules.common.mobileone import reparameterize_model
+from app.core.config import settings
+from app.schemas.staging import StagingRequest
+
+class StagingRanker:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(StagingRanker, cls).__new__(cls)
+            cls._instance.load_model()
+        return cls._instance
+
+    def load_model(self):
+        """
+        Loads the MobileCLIP2 model and tokenizer into memory.
+        This is a time-consuming operation and should only be done once.
+        """
+        print("Loading MobileCLIP2 model...")
+        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+            settings.MODEL_NAME, 
+            pretrained=settings.MODEL_PATH
+        )
+        self.tokenizer = open_clip.get_tokenizer(settings.MODEL_NAME)
+        self.model.eval()
+        self.model = reparameterize_model(self.model)
+        print("Model loaded successfully.")
+
+    def compute_score(self, image: Image.Image, prompts: StagingRequest) -> float:
+        """
+        Computes the differential stageability score for a given image and prompts.
+        """
+        image_tensor = self.preprocess(image.convert("RGB")).unsqueeze(0)
+        
+        text_prompts = [prompts.prompt_good, prompts.prompt_bad]
+        if prompts.prompt_aesthetic:
+            text_prompts.append(prompts.prompt_aesthetic)
+            
+        text_tokens = self.tokenizer(text_prompts)
+
+        with torch.no_grad():
+            image_features = self.model.encode_image(image_tensor)
+            text_features = self.model.encode_text(text_tokens)
+
+            image_features /= image_features.norm(dim=-1, keepdim=True)
+            text_features /= text_features.norm(dim=-1, keepdim=True)
+
+            sims = (image_features @ text_features.T)[0]
+            
+            score = (sims[0] - sims[1]).item()
+            
+            if prompts.prompt_aesthetic:
+                score += sims[2].item()
+                
+            return score
+
+staging_ranker = StagingRanker()

app/schemas/staging.py

@@ -0,0 +1,28 @@
+from pydantic import BaseModel, Field
+from typing import Optional
+
+class StagingRequest(BaseModel):
+    """
+    Pydantic model for user-provided prompts.
+    Users can override the default prompts to tune the scoring logic.
+    """
+    prompt_good: str = Field(
+        "an empty room ideal for virtual staging: large visible floor space, clear walls and corners, windows visible and not blocked, no doorway in the middle, evenly lit with natural light, aesthetically pleasing",
+        description="A descriptive prompt for what makes a room suitable for staging."
+    )
+    prompt_bad: str = Field(
+        "a room that is hard to stage: narrow, cluttered, windows blocked, poor lighting, doorway or obstacles in the center, little open space",
+        description="A descriptive prompt for what makes a room unsuitable for staging."
+    )
+    prompt_aesthetic: Optional[str] = Field(
+        None,
+        description="An optional plus-prompt for aesthetic qualities like 'modern fireplace' or 'hardwood floors'."
+    )
+
+class StagingResponse(BaseModel):
+    """
+    Pydantic model for the API response.
+    """
+    filename: str
+    stageability_score: float = Field(..., description="The calculated differential score (good - bad).")
+    details: str

requirements.txt

@@ -0,0 +1,9 @@
+fastapi
+uvicorn[standard]
+pydantic
+python-multipart
+torch
+Pillow
+
+git+https://github.com/mlfoundations/open_clip.git
+git+https://github.com/apple/ml-mobileclip