Update main.py

main.py

@@ -1,13 +1,23 @@
 import os
+import logging
+import json
+import base64
+from typing import Dict, Any
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
 
 # Set Hugging Face cache directory to /tmp
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 os.environ["TORCH_HOME"] = "/tmp/torch"
 
-from fastapi import FastAPI, File, UploadFile, Form
+from fastapi import FastAPI, Form, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import Response
 import uvicorn
 from PIL import Image
 import io
@@ -18,6 +28,13 @@ from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 import torch
 import cv2
+from dotenv import load_dotenv
+import openai
+import requests
+from io import BytesIO
+
+load_dotenv()
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
 app = FastAPI()
 
@@ -35,34 +52,50 @@ os.makedirs("/tmp/huggingface", exist_ok=True)
 os.makedirs("/tmp/torch", exist_ok=True)
 
 # Load the langSAM model
+logger.info("Loading LangSAM model...")
 langsam_model = LangSAM()
+logger.info("LangSAM model loaded successfully")
 
 # Load SAM2 Model
+logger.info("Loading SAM2 model...")
 sam2_checkpoint = "sam2.1_hiera_small.pt"
 model_cfg = "configs/sam2.1/sam2.1_hiera_s.yaml"
 device = torch.device("cpu")
 
 sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
 predictor = SAM2ImagePredictor(sam2_model)
+logger.info("SAM2 model loaded successfully")
 
 @app.get("/")
 async def root():
     return {"message": "LangSAM API is running!"}
 
-def apply_mask(image, mask):
-    """Overlay mask on image."""
-    mask = mask.astype(np.uint8) * 255  # Convert mask to 0-255 scale
-    mask_colored = np.zeros((*mask.shape, 3), dtype=np.uint8)
-    mask_colored[mask > 0] = [30, 144, 255]  # Blue color for the mask
+def create_mask_overlay(image: np.ndarray, mask: np.ndarray, alpha: float = 0.5) -> np.ndarray:
+    """Create a mask overlay on the original image."""
+    # Create a colored mask (blue color)
+    colored_mask = np.zeros_like(image)
+    colored_mask[mask > 0] = [30, 144, 255]  # Blue color
     
     # Add contour
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    cv2.drawContours(mask_colored, contours, -1, (255, 255, 255), thickness=2)
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    cv2.drawContours(colored_mask, contours, -1, (255, 255, 255), thickness=2)
     
     # Blend with original image
-    overlay = cv2.addWeighted(image, 0.7, mask_colored, 0.3, 0)
+    overlay = cv2.addWeighted(image, 1 - alpha, colored_mask, alpha, 0)
     return overlay
 
+def create_mask_only(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """Create an image showing only the masked region."""
+    # Create a black background
+    result = np.zeros_like(image)
+    # Copy only the masked region
+    result[mask > 0] = image[mask > 0]
+    return result
+
+def image_to_base64(image: np.ndarray) -> str:
+    """Convert numpy array image to base64 string."""
+    _, buffer = cv2.imencode('.png', cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+    return base64.b64encode(buffer).decode('utf-8')
 
 def draw_image(image_rgb, masks, xyxy, probs, labels):
     mask_annotator = sv.MaskAnnotator()
@@ -81,72 +114,134 @@ def draw_image(image_rgb, masks, xyxy, probs, labels):
     annotated_image = mask_annotator.annotate(scene=image_rgb.copy(), detections=detections)
     return annotated_image
 
+def load_image_from_url(url):
+    """Fetch image from URL and load it into memory."""
+    try:
+        logger.info(f"Fetching image from URL: {url}")
+        response = requests.get(url)
+        response.raise_for_status()
+        return Image.open(BytesIO(response.content))
+    except Exception as e:
+        logger.error(f"Error loading image from URL: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Error loading image from URL: {str(e)}")
+
+prompt = """You will be provided with a complete product name, which may contain brand names, extra details, and categories. Your task is to extract only the core product name (apparel or accessory) while removing brand names, categories, and unnecessary words and convert it's meaning to a basic clothing or accessory category.
+
+Examples:
+Beachwood Luxe Paneled Unitard — Girlfriend Collective → Dress
+100 cotton strappy top · Black, White, Red, Peach · T-shirts And Polo Shirts | Massimo Dutti → Shirt
+Wide-leg co-ord trousers with pleats · Green · Dressy | Massimo Dutti → Pants
+BLANKNYC Wide Leg Jean in Radio Star | REVOLVE → Jeans
+
+Basically, you need to convert the product name to a basic clothing or accessory category like Shirt, Pants, Dress, Jeans, etc.
+Now, extract the core product name from the following:
+
+{product_name}"""
+
+@app.post("/openai/chat")
+async def chat(product_name: str = Form(...)):
+    try:
+        logger.info(f"Processing product name: {product_name}")
+        completion = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": prompt.format(product_name=product_name)}],
+        )
+        result = completion.choices[0].message
+        logger.info(f"OpenAI response: {result.content}")
+        return result
+    except Exception as e:
+        logger.error(f"Error in OpenAI chat: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing product name: {str(e)}")
+
 @app.post("/segment/sam2")
 async def segment_image(
-    file: UploadFile = File(...), 
+    image_url: str = Form(...), 
     x: int = Form(...), 
     y: int = Form(...)
 ):
     """Segment image using SAM2 with a single input point."""
-    image_bytes = await file.read()
-    image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    image_array = np.array(image_pil)
-    
-    predictor.set_image(image_array)
-    
-    input_point = np.array([[x, y]])
-    input_label = np.array([1])  # Foreground point
-    
-    # Run SAM2 model
-    masks, scores, logits = predictor.predict(
-        point_coords=input_point,
-        point_labels=input_label,
-        multimask_output=True,
-    )
-
-    # Get top mask
-    top_mask = masks[np.argmax(scores)]
-
-    # Apply mask overlay
-    output_image = apply_mask(image_array, top_mask)
-
-    # Convert to PNG
-    output_pil = Image.fromarray(output_image)
-    img_io = io.BytesIO()
-    output_pil.save(img_io, format="PNG")
-    img_io.seek(0)
-
-    return Response(content=img_io.getvalue(), media_type="image/png")
-
+    try:
+        logger.info(f"Starting SAM2 segmentation for image URL: {image_url}")
+        image_pil = load_image_from_url(image_url)
+        image_array = np.array(image_pil)
+        
+        logger.info("Setting image in SAM2 predictor")
+        predictor.set_image(image_array)
+        
+        input_point = np.array([[x, y]])
+        input_label = np.array([1])  # Foreground point
+        
+        logger.info("Running SAM2 prediction")
+        masks, scores, logits = predictor.predict(
+            point_coords=input_point,
+            point_labels=input_label,
+            multimask_output=True,
+        )
+
+        # Get top mask
+        top_mask = masks[np.argmax(scores)]
+
+        # Create different versions of the result
+        overlay_image = create_mask_overlay(image_array, top_mask)
+        mask_only_image = create_mask_only(image_array, top_mask)
+        
+        # Convert images to base64
+        original_b64 = image_to_base64(image_array)
+        overlay_b64 = image_to_base64(overlay_image)
+        mask_only_b64 = image_to_base64(mask_only_image)
+        
+        # Create response
+        response = {
+            "original": original_b64,
+            "overlay": overlay_b64,
+            "mask_only": mask_only_b64,
+            "score": float(scores[np.argmax(scores)])
+        }
+        
+        logger.info("SAM2 segmentation completed successfully")
+        return response
+    except Exception as e:
+        logger.error(f"Error in SAM2 segmentation: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error in SAM2 segmentation: {str(e)}")
 
 @app.post("/segment/langsam")
-async def segment_image(file: UploadFile = File(...), text_prompt: str = Form(...)):
-    image_bytes = await file.read()
-    image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    
-    # Run segmentation
-    results = langsam_model.predict([image_pil], [text_prompt])
-    
-    # Convert to NumPy array
-    image_array = np.asarray(image_pil)
-    output_image = draw_image(
-        image_array,
-        results[0]["masks"],
-        results[0]["boxes"],
-        results[0]["scores"],
-        results[0]["labels"],
-    )
-    
-    # Convert back to PIL Image
-    output_pil = Image.fromarray(np.uint8(output_image)).convert("RGB")
-    
-    # Save to byte stream
-    img_io = io.BytesIO()
-    output_pil.save(img_io, format="PNG")
-    img_io.seek(0)
-    
-    return Response(content=img_io.getvalue(), media_type="image/png")
-
+async def segment_image(image_url: str = Form(...), text_prompt: str = Form(...)):
+    try:
+        logger.info(f"Starting LangSAM segmentation for image URL: {image_url} with prompt: {text_prompt}")
+        image_pil = load_image_from_url(image_url)
+        image_array = np.array(image_pil)
+        
+        # Run segmentation
+        logger.info("Running LangSAM prediction")
+        results = langsam_model.predict([image_pil], [text_prompt])
+        
+        # Get the first (best) mask
+        mask = results[0]["masks"][0]
+        
+        # Create different versions of the result
+        overlay_image = create_mask_overlay(image_array, mask)
+        mask_only_image = create_mask_only(image_array, mask)
+        
+        # Convert images to base64
+        original_b64 = image_to_base64(image_array)
+        overlay_b64 = image_to_base64(overlay_image)
+        mask_only_b64 = image_to_base64(mask_only_image)
+        
+        # Create response
+        response = {
+            "original": original_b64,
+            "overlay": overlay_b64,
+            "mask_only": mask_only_b64,
+            "boxes": results[0]["boxes"].tolist(),
+            "scores": results[0]["scores"].tolist(),
+            "labels": results[0]["labels"]
+        }
+        
+        logger.info("LangSAM segmentation completed successfully")
+        return response
+    except Exception as e:
+        logger.error(f"Error in LangSAM segmentation: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error in LangSAM segmentation: {str(e)}")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)