updated

app/api/endpoints/staging.py

@@ -1,35 +1,42 @@
 from fastapi import APIRouter, File, UploadFile, Depends, HTTPException
+from typing import List
 from PIL import Image
 import io
 from app.models.clip_model import staging_ranker
-from app.schemas.staging import StagingRequest, StagingResponse
+from app.schemas.staging import StagingRequest, RankingResponse, RankedImage
 
 router = APIRouter()
 
-@router.post("/rank_image", response_model=StagingResponse)
-async def rank_image_for_staging(
+@router.post("/rank_images", response_model=RankingResponse)
+async def rank_images_for_staging(
     prompts: StagingRequest = Depends(),
-    file: UploadFile = File(...)
+    files: List[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`.
+    Accepts multiple image uploads and optional JSON prompts to compute stageability scores.
+    Returns a sorted list of all images and identifies the best one.
     """
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="File provided is not an image.")
+    if not files:
+        raise HTTPException(status_code=400, detail="No files were uploaded.")
+
+    results = []
+    for file in files:
+        if not file.content_type.startswith("image/"):
+            continue
+        
+        try:
+            contents = await file.read()
+            image = Image.open(io.BytesIO(contents))
+            score = staging_ranker.compute_score(image, prompts)
+            results.append(RankedImage(filename=file.filename, stageability_score=score))
+        except Exception as e:
+            continue 
+
+    if not results:
+        raise HTTPException(status_code=400, detail="No valid images were found in the upload.")
 
-    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.")
+    results.sort(key=lambda x: x.stageability_score, reverse=True)
 
-    score = staging_ranker.compute_score(image, prompts)
+    best_image = results[0]
     
-    return StagingResponse(
-        filename=file.filename,
-        stageability_score=score,
-        details="Score calculated based on the provided prompts."
-    )
+    return RankingResponse(best_image=best_image, ranked_list=results)

app/schemas/staging.py

@@ -1,11 +1,7 @@
 from pydantic import BaseModel, Field
-from typing import Optional
+from typing import Optional, List
 
 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."
@@ -19,10 +15,10 @@ class StagingRequest(BaseModel):
         description="An optional plus-prompt for aesthetic qualities like 'modern fireplace' or 'hardwood floors'."
     )
 
-class StagingResponse(BaseModel):
-    """
-    Pydantic model for the API response.
-    """
+class RankedImage(BaseModel):
     filename: str
-    stageability_score: float = Field(..., description="The calculated differential score (good - bad).")
-    details: str
+    stageability_score: float
+
+class RankingResponse(BaseModel):
+    best_image: RankedImage
+    ranked_list: List[RankedImage]

app/ui/gradio_ui.py

@@ -1,16 +1,15 @@
 import gradio as gr
 from PIL import Image
+import json
 from app.models.clip_model import staging_ranker
 from app.schemas.staging import StagingRequest
 
-def rank_image_ui(image, prompt_good, prompt_bad, prompt_aesthetic):
+def rank_images_ui(files, prompt_good, prompt_bad, prompt_aesthetic, progress=gr.Progress(track_tqdm=True)):
     """
-    Bridge function between the Gradio UI and our model engine.
+    Bridge function for the Gradio UI to handle multiple files.
     """
-    if image is None:
-        raise gr.Error("Please upload an image to rank.")
-
-    pil_image = Image.fromarray(image)
+    if files is None:
+        raise gr.Error("Please upload at least one image to rank.")
 
     aesthetic = prompt_aesthetic if prompt_aesthetic.strip() else None
     prompts = StagingRequest(
@@ -19,52 +18,59 @@ def rank_image_ui(image, prompt_good, prompt_bad, prompt_aesthetic):
         prompt_aesthetic=aesthetic
     )
 
-    score = staging_ranker.compute_score(pil_image, prompts)
+    results = []
+    for file in progress.tqdm(files, desc="Ranking Images..."):
+        try:
+            pil_image = Image.open(file.name)
+            score = staging_ranker.compute_score(pil_image, prompts)
+            results.append({"filename": file.name.split('/')[-1], "stageability_score": round(score, 4)})
+        except Exception:
+            continue 
+
+    if not results:
+        raise gr.Error("No valid images could be processed.")
+
+    results.sort(key=lambda x: x["stageability_score"], reverse=True)
 
-    return round(score, 4)
+    best_filename = results[0]["filename"]
+    best_image_path = next((f.name for f in files if f.name.endswith(best_filename)), None)
+
+    return best_image_path, results
 
 
 default_prompts = StagingRequest()
 
-with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 960px !important;}") as gradio_app:
-    gr.Markdown("# 🤖🖼️ Auto Pick MobileCLIP")
-    gr.Markdown("Upload an image and provide text prompts to score its quality and suitability. The higher the score, the better the match.")
+with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 1080px !important;}") as gradio_app:
+    gr.Markdown("# Auto Pick MobileCLIP")
+    gr.Markdown("Upload a batch of images to find the best one. The model scores and ranks all images based on your criteria.")
     
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(type="numpy", label="Upload Image")
+            file_input = gr.File(
+                label="Upload Images",
+                file_count="multiple",
+                file_types=["image"]
+            )
             
             gr.Markdown("### Customize Scoring Prompts")
-            prompt_good_input = gr.Textbox(
-                label="✅ Good Prompt",
-                value=default_prompts.prompt_good,
-                lines=3
-            )
-            prompt_bad_input = gr.Textbox(
-                label="❌ Bad Prompt",
-                value=default_prompts.prompt_bad,
-                lines=3
-            )
-            prompt_aesthetic_input = gr.Textbox(
-                label="✨ Aesthetic 'Plus' Prompt (Optional)",
-                placeholder="e.g., 'features a modern fireplace' or 'hardwood floors'",
-                lines=2
-            )
+            prompt_good_input = gr.Textbox(label="✅ Good Prompt", value=default_prompts.prompt_good, lines=3)
+            prompt_bad_input = gr.Textbox(label="❌ Bad Prompt", value=default_prompts.prompt_bad, lines=3)
+            prompt_aesthetic_input = gr.Textbox(label="✨ Aesthetic 'Plus' Prompt (Optional)", placeholder="e.g., 'features a modern fireplace'", lines=2)
             
-            submit_button = gr.Button("Rank Image", variant="primary")
+            submit_button = gr.Button("Rank Images", variant="primary")
 
-        with gr.Column(scale=1):
-            output_score = gr.Number(label="🏆 Stageability Score", interactive=False)
+        with gr.Column(scale=2):
+            gr.Markdown("### Results")
+            best_image_output = gr.Image(label="🏆 Best Image", interactive=False)
+            full_ranking_output = gr.JSON(label="Full Ranking")
 
     submit_button.click(
-        fn=rank_image_ui,
+        fn=rank_images_ui,
         inputs=[
-            image_input,
+            file_input,
             prompt_good_input,
             prompt_bad_input,
             prompt_aesthetic_input
         ],
-        outputs=output_score
-    )
-
-    
+        outputs=[best_image_output, full_ranking_output]
+    )