from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel import numpy as np import cv2 from PIL import Image import io import base64 from typing import List, Optional import torch from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor import uvicorn app = FastAPI(title="Wall Color Visualizer API") # Configure CORS app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Global variables for SAM model sam_checkpoint = "sam_vit_h_4b8939.pth" model_type = "vit_h" device = "cuda" if torch.cuda.is_available() else "cpu" sam = None mask_generator = None predictor = None # Request models class SegmentRequest(BaseModel): image_base64: str point_x: Optional[float] = None point_y: Optional[float] = None class ColorChangeRequest(BaseModel): image_base64: str mask_base64: str color_hex: str opacity: float = 0.8 # Initialize SAM model def initialize_sam(): global sam, mask_generator, predictor try: print(f"Loading SAM model on {device}...") sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) mask_generator = SamAutomaticMaskGenerator(sam) predictor = SamPredictor(sam) print("SAM model loaded successfully!") except Exception as e: print(f"Warning: Could not load SAM model: {e}") print("The API will run but segmentation features will be limited.") @app.on_event("startup") async def startup_event(): initialize_sam() @app.get("/") async def root(): return { "message": "Wall Color Visualizer API", "status": "running", "sam_loaded": sam is not None } @app.get("/health") async def health_check(): return { "status": "healthy", "device": device, "sam_model_loaded": sam is not None } def decode_base64_image(base64_string: str) -> np.ndarray: """Decode base64 string to numpy array image""" try: # Remove data URL prefix if present if "base64," in base64_string: base64_string = base64_string.split("base64,")[1] img_data = base64.b64decode(base64_string) img = Image.open(io.BytesIO(img_data)) img_array = np.array(img.convert("RGB")) return img_array except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid image data: {str(e)}") def encode_image_to_base64(image: np.ndarray) -> str: """Encode numpy array image to base64 string""" img = Image.fromarray(image.astype(np.uint8)) buffered = io.BytesIO() img.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode() return img_str def encode_mask_to_base64(mask: np.ndarray) -> str: """Encode binary mask to base64 string""" mask_uint8 = (mask * 255).astype(np.uint8) img = Image.fromarray(mask_uint8) buffered = io.BytesIO() img.save(buffered, format="PNG") mask_str = base64.b64encode(buffered.getvalue()).decode() return mask_str def hex_to_rgb(hex_color: str) -> tuple: """Convert hex color to RGB tuple""" hex_color = hex_color.lstrip('#') return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) @app.post("/segment-automatic") async def segment_automatic(file: UploadFile = File(...)): """Automatically segment all objects in the image""" if sam is None: raise HTTPException(status_code=503, detail="SAM model not loaded") try: # Read and decode image contents = await file.read() image = Image.open(io.BytesIO(contents)) image_np = np.array(image.convert("RGB")) # Generate masks masks = mask_generator.generate(image_np) # Sort masks by area (largest first) masks = sorted(masks, key=lambda x: x['area'], reverse=True) # Return top masks result_masks = [] for i, mask_data in enumerate(masks[:2]): # Return top 10 masks mask = mask_data['segmentation'] result_masks.append({ "id": i, "mask_base64": encode_mask_to_base64(mask), "area": int(mask_data['area']), "bbox": [int(x) for x in mask_data['bbox']] }) return { "success": True, "num_masks": len(result_masks), "masks": result_masks, "image_base64": encode_image_to_base64(image_np) } except Exception as e: raise HTTPException(status_code=500, detail=f"Segmentation failed: {str(e)}") @app.post("/segment-point") async def segment_point(request: SegmentRequest): """Segment object at a specific point in the image""" if sam is None: raise HTTPException(status_code=503, detail="SAM model not loaded") try: # Decode image image_np = decode_base64_image(request.image_base64) # Set image for predictor predictor.set_image(image_np) # Use point prompt if request.point_x is not None and request.point_y is not None: point_coords = np.array([[request.point_x, request.point_y]]) point_labels = np.array([1]) # 1 = foreground point masks, scores, logits = predictor.predict( point_coords=point_coords, point_labels=point_labels, multimask_output=True ) # Get the best mask (highest score) best_mask_idx = np.argmax(scores) best_mask = masks[best_mask_idx] return { "success": True, "mask_base64": encode_mask_to_base64(best_mask), "score": float(scores[best_mask_idx]) } else: raise HTTPException(status_code=400, detail="Point coordinates required") except Exception as e: raise HTTPException(status_code=500, detail=f"Segmentation failed: {str(e)}") @app.post("/apply-color") async def apply_color(request: ColorChangeRequest): """Apply color to masked region of the image""" try: # Decode image and mask image_np = decode_base64_image(request.image_base64) mask_np = decode_base64_image(request.mask_base64) # Convert mask to binary if len(mask_np.shape) == 3: mask_np = cv2.cvtColor(mask_np, cv2.COLOR_RGB2GRAY) mask_binary = (mask_np > 128).astype(np.uint8) # Convert hex color to RGB rgb_color = hex_to_rgb(request.color_hex) # Create colored overlay colored_mask = np.zeros_like(image_np) colored_mask[mask_binary == 1] = rgb_color # Blend with original image result = image_np.copy().astype(float) alpha = request.opacity result[mask_binary == 1] = ( alpha * colored_mask[mask_binary == 1] + (1 - alpha) * image_np[mask_binary == 1] ) result = result.astype(np.uint8) return { "success": True, "result_base64": encode_image_to_base64(result) } except Exception as e: raise HTTPException(status_code=500, detail=f"Color application failed: {str(e)}") @app.post("/simple-segment") async def simple_segment(file: UploadFile = File(...)): """Simple segmentation using traditional CV methods (fallback when SAM not available)""" try: # Read and decode image contents = await file.read() image = Image.open(io.BytesIO(contents)) image_np = np.array(image.convert("RGB")) # Convert to different color spaces for better wall detection hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV) gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) # Apply edge detection edges = cv2.Canny(gray, 50, 150) # Dilate edges to create connected regions kernel = np.ones((5, 5), np.uint8) dilated = cv2.dilate(edges, kernel, iterations=2) # Find contours contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Create masks for largest contours result_masks = [] h, w = image_np.shape[:2] # Sort by area contours = sorted(contours, key=cv2.contourArea, reverse=True) for i, contour in enumerate(contours[:5]): # Top 5 regions area = cv2.contourArea(contour) if area < (h * w * 0.01): # Skip very small regions continue mask = np.zeros((h, w), dtype=np.uint8) cv2.drawContours(mask, [contour], -1, 255, -1) # Get bounding box x, y, bw, bh = cv2.boundingRect(contour) result_masks.append({ "id": i, "mask_base64": encode_mask_to_base64(mask / 255), "area": int(area), "bbox": [int(x), int(y), int(bw), int(bh)] }) return { "success": True, "num_masks": len(result_masks), "masks": result_masks, "image_base64": encode_image_to_base64(image_np), "method": "traditional_cv" } except Exception as e: raise HTTPException(status_code=500, detail=f"Segmentation failed: {str(e)}") if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)