Week 3 complete: canvas, drag, resize, delete, export, history

config.py

@@ -4,8 +4,9 @@ from dotenv import load_dotenv
 load_dotenv()
 
 PROJECT_NAME = os.getenv("PROJECT_NAME", "poster-editor")
-UPLOAD_DIR = os.getenv("UPLOAD_DIR", "uploads")
-OUTPUT_DIR = os.getenv("OUTPUT_DIR", "outputs")
+BACKEND_DIR = os.path.dirname(os.path.abspath(__file__))
+UPLOAD_DIR  = os.path.join(BACKEND_DIR, os.getenv("UPLOAD_DIR", "uploads"))
+OUTPUT_DIR  = os.path.join(BACKEND_DIR, os.getenv("OUTPUT_DIR", "outputs"))
 MAX_FILE_SIZE_MB = int(os.getenv("MAX_FILE_SIZE_MB", 20))
 
 os.makedirs(UPLOAD_DIR, exist_ok=True)

grounding.py

@@ -52,8 +52,6 @@ def detect_objects(image_path: str, model, text_prompt: str = None) -> list:
     # load_image returns (PIL image, transformed tensor) — both needed
     image_pil, image_tensor = load_image(image_path)
 
-    print(type(image_pil))
-    print(image_pil.shape)
 
     img_h, img_w = image_pil.shape[:2]
 

layer_extractor.py

@@ -1,49 +1,120 @@
 import os
+import sys
 import json
+import base64
+import io
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+
 from grounding import load_grounding_model, detect_objects
 from ocr import extract_text
 
 BACKEND_DIR = os.path.dirname(os.path.abspath(__file__))
-
-# labels that GroundingDINO detects but we handle separately or don't need
 SKIP_LABELS = {"background", "text", "button"}
 
 
+def pil_to_base64(img: Image.Image, fmt: str = "PNG") -> str:
+    """Convert PIL image to base64 string."""
+    buffer = io.BytesIO()
+    img.save(buffer, format=fmt)
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def crop_text_layer(image_path: str, x: int, y: int, w: int, h: int) -> str:
+    """
+    Crop text region as transparent PNG.
+    Text layers use simple rectangular crop with white pixels made transparent.
+    """
+    img  = Image.open(image_path).convert("RGBA")
+    crop = img.crop((x, y, x + w, y + h))
+    return pil_to_base64(crop, "PNG")
+
+
 def build_layers(image_path: str) -> list:
     """
-    Run GroundingDINO + OCR on image, combine results into a flat layer list.
-    Each layer has: id, type, x, y, w, h, confidence (+ text for text layers).
-    This JSON is what the Fabric.js canvas loads to render editable elements.
+    Full layer extraction pipeline:
+    1. GroundingDINO finds named objects with bounding boxes
+    2. SAM2 refines each box into a precise pixel mask
+    3. PaddleOCR finds all text blocks
+    4. Each layer gets a transparent PNG crop
+
+    Returns flat list of layers ready for Fabric.js canvas.
     """
     layers   = []
     layer_id = 1
 
-    # --- object layers via GroundingDINO ---
-    print("[LAYERS] Running object detection...")
+    # --- load image ---
+    image_bgr = cv2.imread(image_path)
+    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    img_h, img_w = image_rgb.shape[:2]
+
+    # --- step 1: GroundingDINO object detection ---
+    print("[LAYERS] Running GroundingDINO...")
     gdino_model = load_grounding_model()
     detections  = detect_objects(image_path, gdino_model)
 
-    for det in detections:
-        if det["label"].lower() in SKIP_LABELS:
-            continue
-
-        layers.append({
-            "id":         layer_id,
-            "type":       "object",
-            "label":      det["label"],
-            "x":          det["x1"],
-            "y":          det["y1"],
-            "w":          det["w"],
-            "h":          det["h"],
-            "confidence": det["confidence"],
-        })
-        layer_id += 1
-
-    # --- text layers via PaddleOCR ---
+    # filter background/text labels
+    obj_detections = [d for d in detections if d["label"].lower() not in SKIP_LABELS]
+    print(f"[LAYERS] {len(obj_detections)} objects to segment with SAM2")
+
+    # --- step 2: SAM2 precise masks for each object ---
+    if obj_detections:
+        from segment import load_sam2_model, get_mask_for_box, mask_to_transparent_png
+        predictor = load_sam2_model()
+
+        for det in obj_detections:
+            box = [det["x1"], det["y1"], det["x2"], det["y2"]]
+
+            try:
+                mask = get_mask_for_box(predictor, image_rgb, box)
+
+                # get tight bounding box from mask
+                rows     = np.where(mask.any(axis=1))[0]
+                cols     = np.where(mask.any(axis=0))[0]
+                if len(rows) == 0: continue
+
+                y1, y2 = int(rows.min()), int(rows.max())
+                x1, x2 = int(cols.min()), int(cols.max())
+
+                # create transparent PNG
+                png_img = mask_to_transparent_png(image_rgb, mask)
+                b64     = pil_to_base64(png_img, "PNG")
+
+                layers.append({
+                    "id":         layer_id,
+                    "type":       "object",
+                    "label":      det["label"],
+                    "x":          x1,
+                    "y":          y1,
+                    "w":          x2 - x1,
+                    "h":          y2 - y1,
+                    "confidence": det["confidence"],
+                    "base64":     b64,
+                    "format":     "png",
+                })
+                layer_id += 1
+
+            except Exception as e:
+                print(f"[LAYERS] SAM2 failed for {det['label']}: {e}")
+                continue
+
+        # free SAM2 VRAM before OCR
+        del predictor
+        torch.cuda.empty_cache() if __import__('torch').cuda.is_available() else None
+        print("[LAYERS] SAM2 VRAM freed")
+
+    # --- step 3: OCR text layers ---
     print("[LAYERS] Running OCR...")
     text_blocks = extract_text(image_path)
 
     for block in text_blocks:
+        b64 = crop_text_layer(
+            image_path,
+            block["x"], block["y"],
+            block["w"], block["h"]
+        )
         layers.append({
             "id":         layer_id,
             "type":       "text",
@@ -53,10 +124,12 @@ def build_layers(image_path: str) -> list:
             "w":          block["w"],
             "h":          block["h"],
             "confidence": block["confidence"],
+            "base64":     b64,
+            "format":     "png",
         })
         layer_id += 1
 
-    print(f"[LAYERS] Built {len(layers)} layers ({len(detections)} objects + {len(text_blocks)} text)")
+    print(f"[LAYERS] Built {len(layers)} layers ({len(obj_detections)} objects + {len(text_blocks)} text)")
     return layers
 
 
@@ -81,20 +154,18 @@ if __name__ == "__main__":
         exit(1)
 
     print("=" * 50)
-    print("Building Editable Layers")
+    print("Building Transparent Layers")
     print("=" * 50)
 
     layers = build_layers(TEST_IMAGE)
 
     print(f"\nGenerated {len(layers)} layers:")
     for layer in layers:
+        fmt = layer.get("format", "jpg")
         if layer["type"] == "text":
-            print(f"  [{layer['id']}] TEXT     '{layer['text']}' at ({layer['x']},{layer['y']})")
+            print(f"  [{layer['id']}] TEXT   '{layer['text']}' at ({layer['x']},{layer['y']}) [{fmt}]")
         else:
-            print(f"  [{layer['id']}] OBJECT   '{layer['label']}' at ({layer['x']},{layer['y']})")
+            print(f"  [{layer['id']}] OBJECT '{layer['label']}' at ({layer['x']},{layer['y']}) [{fmt}]")
 
     save_layers(layers)
-
-    print("=" * 50)
-    print("Open outputs/layers.json")
     print("=" * 50)

main.py

@@ -76,14 +76,15 @@ async def upload_image(
         raise HTTPException(500, f"Pipeline failed: {str(e)}")
 
     return JSONResponse({
-        "project_id":        project.id,
-        "file_id":           file_id,
-        "filename":          file.filename,
-        "image_w":           result["image_w"],
-        "image_h":           result["image_h"],
-        "image_base64":      result["image_base64"],
-        "layers":            result["layers"],
-        "processing_time_s": result["processing_time_s"],
+        "project_id":           project.id,
+        "file_id":              file_id,
+        "filename":             file.filename,
+        "image_w":              result["image_w"],
+        "image_h":              result["image_h"],
+        "background_base64":    result["background_base64"],
+        "original_base64":      result["original_base64"],
+        "layers":               result["layers"],
+        "processing_time_s":    result["processing_time_s"],
     })
 
 
@@ -166,4 +167,31 @@ async def inpaint_element(request: dict):
 
     except Exception as e:
         unload_inpaint_model()
-        raise HTTPException(500, f"Inpainting failed: {str(e)}")
+        raise HTTPException(500, f"Inpainting failed: {str(e)}")
+    
+
+@app.get("/projects/{file_id}/layers")
+async def get_project_layers(file_id: str, db: Session = Depends(get_db)):
+    """Re-run pipeline on existing upload to get full layer data."""
+    project = get_project_by_file_id(db, file_id)
+    if not project:
+        raise HTTPException(404, "Project not found")
+
+    if not os.path.exists(project.upload_path):
+        raise HTTPException(404, "Original file no longer exists")
+
+    try:
+        result = run_pipeline(project.upload_path)
+        return JSONResponse({
+            "project_id":        project.id,
+            "file_id":           file_id,
+            "filename":          project.filename,
+            "image_w":           result["image_w"],
+            "image_h":           result["image_h"],
+            "background_base64": result["background_base64"],
+            "original_base64":   result["original_base64"],
+            "layers":            result["layers"],
+            "processing_time_s": result["processing_time_s"],
+        })
+    except Exception as e:
+        raise HTTPException(500, f"Pipeline failed: {str(e)}")

pipeline.py

@@ -2,98 +2,105 @@ import os
 import time
 import base64
 import json
+import io
+import cv2
+import numpy as np
 from PIL import Image
 
-from layer_extractor import build_layers, save_layers
+from layer_extractor import build_layers
 
 BACKEND_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
 def image_to_base64(image_path: str) -> str:
-    """Read image file and return base64 string for sending to frontend."""
     with open(image_path, "rb") as f:
         return base64.b64encode(f.read()).decode("utf-8")
 
 
-def crop_layer_image(image_path: str, x: int, y: int, w: int, h: int) -> str:
+def pil_to_base64(img: Image.Image, fmt: str = "PNG") -> str:
+    buffer = io.BytesIO()
+    img.save(buffer, format=fmt)
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def reconstruct_background(image_path: str, layers: list) -> str:
     """
-    Crop a region from the original image and return as base64.
-    This gives the frontend the actual pixel content of each layer
-    so Fabric.js can render it on the canvas.
+    Build ONE combined mask covering all detected layers.
+    Run a single OpenCV TELEA inpaint to reconstruct the clean background.
+    Returns base64 JPEG of the clean background.
     """
-    img    = Image.open(image_path).convert("RGB")
-    region = img.crop((x, y, x + w, y + h))
+    image_bgr = cv2.imread(image_path)
+    h, w      = image_bgr.shape[:2]
+
+    # combined mask — white where any layer exists
+    combined_mask = np.zeros((h, w), dtype=np.uint8)
 
-    # save cropped region to temp file then encode
-    temp_path = os.path.join(BACKEND_DIR, "outputs", f"crop_temp.jpg")
-    region.save(temp_path, quality=95)
+    for layer in layers:
+        x  = max(0, layer["x"])
+        y  = max(0, layer["y"])
+        x2 = min(w, layer["x"] + layer["w"])
+        y2 = min(h, layer["y"] + layer["h"])
+        combined_mask[y:y2, x:x2] = 255
+
+    # single TELEA inpaint — fast, works for MVP, replaceable later
+    clean_bg = cv2.inpaint(
+        image_bgr,
+        combined_mask,
+        inpaintRadius=25,
+        flags=cv2.INPAINT_TELEA
+    )
 
-    return image_to_base64(temp_path)
+    # encode to base64
+    _, buffer = cv2.imencode('.jpg', clean_bg, [cv2.IMWRITE_JPEG_QUALITY, 95])
+    return base64.b64encode(buffer).decode("utf-8")
 
 
 def run_pipeline(image_path: str) -> dict:
     """
-    Full pipeline: image → layer JSON ready for Fabric.js canvas.
-
-    Steps:
-    1. GroundingDINO detects named objects
-    2. PaddleOCR reads all text
-    3. Each layer gets its cropped image as base64
-    4. Returns complete response the frontend expects
-
-    Returns:
-        {
-            image_w, image_h,
-            image_base64,        ← full poster for canvas background
-            layers: [
-                {id, type, label/text, x, y, w, h, confidence, base64}
-            ]
-        }
+    Full pipeline:
+    1. Extract all layers as transparent PNGs
+    2. Build combined mask of all layer regions
+    3. Reconstruct clean background in ONE inpaint pass
+    4. Return clean background + independent layers
+
+    Editor renders: clean background + layer PNGs
+    Every element exists exactly once — no duplicates possible.
+    Editing is instant — no AI calls needed during editing.
     """
     start = time.time()
-    print(f"\n[PIPELINE] Starting pipeline for: {os.path.basename(image_path)}")
+    print(f"\n[PIPELINE] Starting: {os.path.basename(image_path)}")
 
-    # get image dimensions
-    img    = Image.open(image_path)
+    img      = Image.open(image_path)
     img_w, img_h = img.size
-    print(f"[PIPELINE] Image size: {img_w}x{img_h}")
+    print(f"[PIPELINE] Image: {img_w}x{img_h}")
 
-    # run detection + OCR
+    # step 1 — extract all layers
     layers = build_layers(image_path)
+    print(f"[PIPELINE] Extracted {len(layers)} layers")
 
-    # attach cropped base64 image to each layer
-    # frontend needs this to render each element on canvas
-    print("[PIPELINE] Cropping layer images...")
-    for layer in layers:
-        try:
-            layer["base64"] = crop_layer_image(
-                image_path,
-                layer["x"], layer["y"],
-                layer["w"], layer["h"]
-            )
-        except Exception as e:
-            print(f"[PIPELINE] Crop failed for layer {layer['id']}: {e}")
-            layer["base64"] = ""
+    # step 2 — reconstruct background once using combined mask
+    print("[PIPELINE] Reconstructing clean background...")
+    bg_base64 = reconstruct_background(image_path, layers)
+    print("[PIPELINE] Background reconstructed")
 
     elapsed = round(time.time() - start, 2)
-    print(f"[PIPELINE] Done in {elapsed}s — {len(layers)} layers extracted")
+    print(f"[PIPELINE] Done in {elapsed}s")
 
     return {
-        "image_w":      img_w,
-        "image_h":      img_h,
-        "image_base64": image_to_base64(image_path),
-        "layers":       layers,
-        "processing_time_s": elapsed,
+        "image_w":              img_w,
+        "image_h":              img_h,
+        "background_base64":    bg_base64,     # clean background, no elements
+        "original_base64":      image_to_base64(image_path),  # for inpainting reference
+        "layers":               layers,
+        "processing_time_s":    elapsed,
     }
 
 
 def save_pipeline_output(result: dict, output_dir: str = None) -> str:
-    """Save pipeline result to JSON — useful for debugging and development."""
     if output_dir is None:
         output_dir = os.path.join(BACKEND_DIR, "outputs")
     os.makedirs(output_dir, exist_ok=True)
 
-    # save without base64 blobs — they're huge and unreadable in JSON viewer
     result_slim = {
         "image_w":           result["image_w"],
         "image_h":           result["image_h"],
@@ -109,7 +116,7 @@ def save_pipeline_output(result: dict, output_dir: str = None) -> str:
     with open(out_path, "w", encoding="utf-8") as f:
         json.dump(result_slim, f, indent=2, ensure_ascii=False)
 
-    print(f"[PIPELINE] Saved slim output → {out_path}")
+    print(f"[PIPELINE] Saved → {out_path}")
     return out_path
 
 
@@ -121,16 +128,14 @@ if __name__ == "__main__":
         exit(1)
 
     print("=" * 50)
-    print("Running Full Editify Pipeline")
+    print("Running Editify Pipeline")
     print("=" * 50)
 
     result = run_pipeline(TEST_IMAGE)
     save_pipeline_output(result)
 
     print("=" * 50)
-    print(f"Pipeline complete.")
-    print(f"  Layers:          {len(result['layers'])}")
-    print(f"  Image size:      {result['image_w']}x{result['image_h']}")
-    print(f"  Processing time: {result['processing_time_s']}s")
-    print(f"  Image base64:    {len(result['image_base64'])} chars")
+    print(f"Layers:     {len(result['layers'])}")
+    print(f"Image:      {result['image_w']}x{result['image_h']}")
+    print(f"Time:       {result['processing_time_s']}s")
     print("=" * 50)

segment.py

@@ -1,27 +1,23 @@
-import torch
-import numpy as np
-import cv2
 import os
-import json
 import sys
+import torch
+import numpy as np
+from PIL import Image
 
-# tell Python where SAM2 code lives
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'sam2'))
 
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-# config
 MODEL_CFG  = "configs/sam2.1/sam2.1_hiera_s.yaml"
 CHECKPOINT = os.path.join(os.path.dirname(__file__), "models", "sam2.1_hiera_small.pt")
 DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
-SAM2_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sam2')
+SAM2_DIR   = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sam2')
 
 print(f"[SAM2] Using device: {DEVICE}")
 
 
-def load_model():
-    """Load SAM2 into GPU memory. Call once at server startup."""
+def load_sam2_model():
     os.chdir(SAM2_DIR)
     model     = build_sam2(MODEL_CFG, CHECKPOINT, device=DEVICE)
     predictor = SAM2ImagePredictor(model)
@@ -29,197 +25,50 @@ def load_model():
     return predictor
 
 
-def generate_grid_points(image_h, image_w, rows=8, cols=8):
+def get_mask_for_box(predictor, image_rgb: np.ndarray, box: list) -> np.ndarray:
     """
-    Create a grid of points spread across the image.
-    SAM2 needs input points as prompts to know where to look.
-    A 5x5 grid = 25 points covering the entire poster evenly.
+    Get precise pixel mask for a single bounding box using SAM2.
+    box = [x1, y1, x2, y2] in absolute pixels.
+    Returns boolean mask same size as image.
     """
-    points = []
-    for r in range(1, rows + 1):
-        for c in range(1, cols + 1):
-            x = int(image_w * c / (cols + 1))
-            y = int(image_h * r / (rows + 1))
-            points.append([x, y])
-    return np.array(points, dtype=np.float32)
-
-
-def segment_image(image_path: str, predictor) -> dict:
-    """
-    Main segmentation function.
-    Input : path to any poster image
-    Output: dict with list of layers [{id, score, x, y, w, h, area, mask}]
-    """
-    # load image
-    image_bgr = cv2.imread(image_path)
-    if image_bgr is None:
-        raise ValueError(f"Could not read image at {image_path}")
-
-    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-    h, w      = image_rgb.shape[:2]
-    print(f"[SAM2] Image loaded: {w}x{h}px")
-
-    # encode image into SAM2 — this is the heavy step (~3-5 seconds)
     predictor.set_image(image_rgb)
-    print("[SAM2] Image encoded, running segmentation...")
 
-    # generate grid prompts
-    grid_points = generate_grid_points(h, w, rows=5, cols=5)
-    labels      = np.ones(len(grid_points), dtype=np.int32)
+    box_array = np.array(box, dtype=np.float32)
 
-    # run SAM2
     with torch.inference_mode():
         with torch.autocast(device_type=DEVICE, dtype=torch.float16):
             masks, scores, _ = predictor.predict(
-                point_coords=grid_points,
-                point_labels=labels,
-                multimask_output=True,
+                point_coords = None,
+                point_labels = None,
+                box          = box_array[None, :],  # SAM2 expects (1, 4)
+                multimask_output = True,
             )
 
-    print(f"[SAM2] Generated {len(masks)} raw masks")
-
-    # flatten masks and scores if shape is (N, H, W)
-    if masks.ndim == 3:
-        masks_list  = [masks[i] for i in range(masks.shape[0])]
-        scores_list = [scores[i] for i in range(scores.shape[0])]
-    else:
-        masks_list  = list(masks)
-        scores_list = list(scores)
-
-    # filter and clean masks
-    layers     = []
-    seen_areas = set()
-
-    for i, (mask, score) in enumerate(zip(masks_list, scores_list)):
-        score = float(score)
-        if score < 0.6:
-            continue
-
-        # get bounding box
-        rows_idx = np.where(mask.any(axis=1))[0]
-        cols_idx = np.where(mask.any(axis=0))[0]
-        if len(rows_idx) == 0 or len(cols_idx) == 0:
-            continue
-
-        y1, y2 = int(rows_idx.min()), int(rows_idx.max())
-        x1, x2 = int(cols_idx.min()), int(cols_idx.max())
-        area    = int((x2 - x1) * (y2 - y1))
-
-        # skip noise and full-image background
-        min_area = int(w * h * 0.005)
-        max_area = int(w * h * 0.95)
-        if area < min_area or area > max_area:
-            continue
-
-        # deduplicate similar masks
-        area_key = area // 1000
-        if area_key in seen_areas:
-            continue
-        seen_areas.add(area_key)
-
-        layers.append({
-            "id":    i,
-            "score": round(score, 3),
-            "x":     x1,
-            "y":     y1,
-            "w":     x2 - x1,
-            "h":     y2 - y1,
-            "area":  area,
-            "mask":  mask.astype(bool),
-        })
-
-    print(f"[SAM2] {len(layers)} clean layers after filtering")
-    return {"layers": layers, "image_size": {"w": w, "h": h}}
-
-
-def save_debug_output(image_path: str, result: dict, output_dir: str = None):
-    if output_dir is None:
-        output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
+    # pick highest confidence mask
+    best_idx  = scores.argmax()
+    best_mask = masks[best_idx].astype(bool)
+    return best_mask
+
+
+def mask_to_transparent_png(image_rgb: np.ndarray, mask: np.ndarray) -> Image.Image:
     """
-    Save visual output so you can SEE what SAM2 detected.
-    - masked_overlay.jpg : poster with coloured regions over each object
-    - layers_data.json   : layer positions and scores
+    Apply mask to image — keep masked pixels, make everything else transparent.
+    Returns RGBA PIL image.
     """
-    os.makedirs(output_dir, exist_ok=True)
-    image_bgr = cv2.imread(image_path)
-    overlay   = image_bgr.copy()
-
-    colours = [
-        (255, 80,  80),  (80,  255, 80),  (80,  80,  255),
-        (255, 255, 80),  (255, 80,  255), (80,  255, 255),
-        (200, 140, 80),  (140, 200, 80),  (80,  140, 200),
-    ]
-
-    json_layers = []
-
-    for idx, layer in enumerate(result["layers"]):
-        colour = colours[idx % len(colours)]
-        mask   = layer["mask"]
-
-        # semi-transparent fill
-        overlay[mask] = (
-            overlay[mask] * 0.45 + np.array(colour) * 0.55
-        ).astype(np.uint8)
-
-        # bounding box
-        cv2.rectangle(
-            overlay,
-            (layer["x"], layer["y"]),
-            (layer["x"] + layer["w"], layer["y"] + layer["h"]),
-            colour, 2
-        )
-
-        # label
-        cv2.putText(
-            overlay,
-            f"L{idx} {layer['score']}",
-            (layer["x"] + 4, layer["y"] + 18),
-            cv2.FONT_HERSHEY_SIMPLEX, 0.5, colour, 1
-        )
-
-        json_layers.append({
-            "id":    layer["id"],
-            "score": layer["score"],
-            "x":     layer["x"],
-            "y":     layer["y"],
-            "w":     layer["w"],
-            "h":     layer["h"],
-            "area":  layer["area"],
-        })
-
-    out_img  = os.path.join(output_dir, "masked_overlay.jpg")
-    out_json = os.path.join(output_dir, "layers_data.json")
-
-    cv2.imwrite(out_img, overlay)
-    with open(out_json, "w") as f:
-        json.dump({
-            "layers":     json_layers,
-            "image_size": result["image_size"]
-        }, f, indent=2)
-
-    print(f"[SAM2] Saved overlay → {out_img}")
-    print(f"[SAM2] Saved JSON   → {out_json}")
-    return out_img, out_json
-
-
-# run directly to test
-if __name__ == "__main__":
-    BACKEND_DIR = os.path.dirname(os.path.abspath(__file__))
-    TEST_IMAGE  = os.path.join(BACKEND_DIR, "test_images", "sale_img.jpg")
-
-    if not os.path.exists(TEST_IMAGE):
-        print(f"ERROR: No image at {TEST_IMAGE}")
-        exit(1)
-
-    print("=" * 50)
-    print("Testing SAM2 segmentation")
-    print("=" * 50)
-
-    predictor = load_model()
-    result    = segment_image(TEST_IMAGE, predictor)
-    save_debug_output(TEST_IMAGE, result)
-
-    print("=" * 50)
-    print(f"DONE. Found {len(result['layers'])} layers.")
-    print("Open outputs/masked_overlay.jpg to see results.")
-    print("=" * 50)
+    rgba        = np.zeros((*image_rgb.shape[:2], 4), dtype=np.uint8)
+    rgba[..., :3] = image_rgb
+    rgba[..., 3]  = (mask * 255).astype(np.uint8)  # alpha = 255 where object is
+
+    # crop to bounding box of mask to minimise image size
+    rows = np.where(mask.any(axis=1))[0]
+    cols = np.where(mask.any(axis=0))[0]
+
+    if len(rows) == 0 or len(cols) == 0:
+        return Image.fromarray(rgba, 'RGBA')
+
+    y1, y2 = int(rows.min()), int(rows.max())
+    x1, x2 = int(cols.min()), int(cols.max())
+
+    cropped = rgba[y1:y2+1, x1:x2+1]
+    return Image.fromarray(cropped, 'RGBA')
+