Add picture preprocessing

app.py

@@ -1,39 +1,68 @@
 import torch
 import torch.nn.functional as F
-from transformers import AutoTokenizer, AutoModel, AutoImageProcessor
+from transformers import AutoModel, AutoImageProcessor
+from PIL import Image
+from rembg import remove
 import gradio as gr
 import spaces
+import io
 
+# Load the Nomic embed model
 processor = AutoImageProcessor.from_pretrained("nomic-ai/nomic-embed-vision-v1.5")
 vision_model = AutoModel.from_pretrained("nomic-ai/nomic-embed-vision-v1.5", trust_remote_code=True)
 
-def ImgEmbed(image):
+def focus_on_subject(image: Image.Image) -> Image.Image:
     """
-    Generate normalized embedding vector for the uploaded image.
+    Remove background and crop to the main object using rembg.
 
     Args:
-        image (PIL.Image.Image or np.ndarray): Input image uploaded by the user.
+        image (PIL.Image.Image): Input image.
 
     Returns:
-        list[float]: A normalized image embedding vector representing the input image.
+        PIL.Image.Image: Cropped image with background removed.
     """
-    print(image);
-    inputs = processor(image, return_tensors="pt")
+    image = image.convert("RGB")
+    
+    # Remove background
+    img_bytes = io.BytesIO()
+    image.save(img_bytes, format="PNG")
+    img_bytes = img_bytes.getvalue()
+    result_bytes = remove(img_bytes)
+    
+    result_image = Image.open(io.BytesIO(result_bytes)).convert("RGBA")
+    bbox = result_image.getbbox()
+    cropped = result_image.crop(bbox) if bbox else result_image
+    
+    return cropped.convert("RGB")
 
+def ImgEmbed(image: Image.Image):
+    """
+    Preprocess image, generate normalized embedding, and return both embedding and processed image.
+
+    Args:
+        image (PIL.Image.Image): Input image.
+
+    Returns:
+        Tuple: (embedding list, processed image)
+    """
+    focused_image = focus_on_subject(image)
+    inputs = processor(focused_image, return_tensors="pt")
     img_emb = vision_model(**inputs).last_hidden_state
     img_embeddings = F.normalize(img_emb[:, 0], p=2, dim=1)
 
-    return img_embeddings[0].tolist();
+    return img_embeddings[0].tolist(), focused_image
 
+# Gradio UI
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            img = gr.Image(label="Upload Image")
+            btn = gr.Button("Get Embeddings")
+        with gr.Column():
+            pre_img = gr.Image(label="Preprocessed Image")
+            out = gr.Text(label="Image Embedding")
 
+    btn.click(ImgEmbed, inputs=[img], outputs=[out, pre_img])
 
-with gr.Blocks() as demo:      
-       img = gr.Image();
-       out = gr.Text();
-       
-       btn = gr.Button("Get Embeddings")
-       btn.click(ImgEmbed, [img], [out])
-       
-       
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch()