Augmentus-robotics/object-analysis-api
0
1from flask import Flask, request, jsonify2import cv23import numpy as np4from PIL import Image, ImageOps5import base646import io7from gradio_client import Client8import json9import os10from functools import wraps11from PIL import Image as PILImage12from utils.segment_model import segment_image13 14app = Flask(__name__)15 16def detect_color_from_hex(image, hex_color, tolerance=40):17 """18 Detects a given hex color in an image and outputs a binary mask [0, 255].19 20 Args:21 image (np.ndarray): Input image in BGR or RGB format.22 hex_color (str): Color in hex format, e.g. '#0000FF' or '0000FF'.23 tolerance (int): Range for color similarity.24 25 Returns:26 np.ndarray: Binary mask (0 or 255).27 """28 # Convert hex to RGB29 hex_color = hex_color.lstrip('#')30 rgb_color = np.array([int(hex_color[i:i+2], 16) for i in (0, 2, 4)])31 32 # Convert RGB to BGR (OpenCV uses BGR)33 target_color = rgb_color[::-1]34 35 # Define bounds with tolerance36 lower_bound = np.array([100, 0, 0]) 37 upper_bound = np.array([255, 100, 100])38 39 lower_bound = np.array([200, 200, 200]) # lower limit for white40 upper_bound = np.array([255, 255, 255]) # pure white41 42 # Create mask43 mask = cv2.inRange(image, lower_bound, upper_bound)44 45 return mask46 47 48def detect_blobs_internal(image, hex_color='#FFFFFF', tolerance=40, min_area=100):49 """50 Robust blob detection using connected components analysis.51 More reliable than SimpleBlobDetector for color-based detection.52 """53 try:54 # --- Image preprocessing ---55 img_array = np.array(image)56 57 if len(img_array.shape) == 3:58 img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)59 else:60 img_bgr = cv2.cvtColor(img_array, cv2.COLOR_GRAY2BGR)61 62 # --- Color mask ---63 binary_mask = detect_color_from_hex(img_bgr, hex_color, tolerance)64 if binary_mask.ndim == 3:65 binary_mask = cv2.cvtColor(binary_mask, cv2.COLOR_BGR2GRAY)66 binary_mask = binary_mask.astype(np.uint8)67 68 # --- Morphological operations (CRITICAL for separating merged blobs) ---69 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))70 # Remove noise71 binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel, iterations=2)72 # Fill small holes73 binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel, iterations=1)74 75 # --- METHOD 1: Connected Components (Fastest) ---76 num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(77 binary_mask, connectivity=8, ltype=cv2.CV_32S78 )79 80 blobs = []81 overlay = img_bgr.copy()82 blob_id = 083 84 # Start from 1 (skip background label 0)85 for i in range(1, num_labels):86 area = float(stats[i, cv2.CC_STAT_AREA])87 88 # Filter by area89 if area < min_area:90 continue91 92 blob_id += 193 94 # Extract stats95 x = int(stats[i, cv2.CC_STAT_LEFT])96 y = int(stats[i, cv2.CC_STAT_TOP])97 w = int(stats[i, cv2.CC_STAT_WIDTH])98 h = int(stats[i, cv2.CC_STAT_HEIGHT])99 cx, cy = centroids[i]100 101 # Extract contour for this component102 component_mask = (labels == i).astype(np.uint8) * 255103 cnt_ret = cv2.findContours(component_mask, cv2.RETR_EXTERNAL, 104 cv2.CHAIN_APPROX_SIMPLE)105 contours = cnt_ret[0] if len(cnt_ret) == 2 else cnt_ret[1]106 107 if len(contours) > 0:108 contour = max(contours, key=cv2.contourArea) # largest contour109 contour_list = contour.reshape(-1, 2).astype(int).tolist()110 else:111 contour_list = []112 contour = None113 114 # Approximate size (diameter of equivalent circle)115 size = 2 * np.sqrt(area / np.pi)116 117 blobs.append({118 "id": blob_id,119 "x": round(float(cx), 2),120 "y": round(float(cy), 2),121 "size": round(size, 2),122 "area": round(area, 2),123 "bbox": [x, y, x + w, y + h],124 "contour": contour_list125 })126 127 # Visualization128 if contour is not None:129 cv2.drawContours(overlay, [contour], -1, (0, 255, 0), 2)130 cv2.circle(overlay, (int(cx), int(cy)), 3, (0, 0, 255), -1)131 cv2.putText(overlay, f"#{blob_id}", (int(cx) + 10, int(cy) - 10),132 cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2, cv2.LINE_AA)133 134 # Header text135 cv2.putText(overlay, f"Objects: {len(blobs)}", (10, 30),136 cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2, cv2.LINE_AA)137 cv2.putText(overlay, f"Color: {hex_color.upper()}", (10, 60),138 cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2, cv2.LINE_AA)139 140 # Encode141 _, buffer = cv2.imencode('.png', overlay)142 img_base64 = base64.b64encode(buffer).decode('utf-8')143 144 return {145 "success": True,146 "num_objects": len(blobs),147 "blobs": blobs,148 "image_shape": {"width": int(img_array.shape[1]), "height": int(img_array.shape[0])},149 "color_detected": hex_color,150 "tolerance_used": tolerance,151 "visualization": img_base64152 }153 154 except Exception as e:155 return {"success": False, "error": str(e), "num_objects": 0, "blobs": []}156 157# --- CROPPER: center-crop to square and resize ---158def center_crop_to_square(image: Image.Image, output_size: int = 1024) -> Image.Image:159 """160 Center-crops any image to a square and resizes to output_size x output_size.161 Accepts a PIL Image and returns a PIL Image.162 """163 164 image = ImageOps.exif_transpose(image)165 166 w, h = image.size167 min_side = min(w, h)168 left = (w - min_side) // 2169 top = (h - min_side) // 2170 right = left + min_side171 bottom = top + min_side172 cropped = image.crop((left, top, right, bottom))173 174 # Pillow 10 renamed resampling enums; keep compatibility175 try:176 resample = Image.Resampling.LANCZOS177 except AttributeError:178 resample = Image.LANCZOS179 180 return cropped.resize((output_size, output_size), resample)181 182 183# --- Helpers (put near your other imports / helpers) ---184def _decode_base64_to_pil(s: str) -> Image.Image:185 if isinstance(s, str) and s.startswith('data:image'):186 s = s.split(',', 1)[1]187 img_bytes = base64.b64decode(s)188 return Image.open(io.BytesIO(img_bytes))189 190def _pil_to_base64_png(img: Image.Image) -> str:191 buf = io.BytesIO()192 img.save(buf, format='PNG')193 buf.seek(0)194 return base64.b64encode(buf.read()).decode('utf-8')195 196def _ensure_lanczos():197 try:198 return Image.Resampling.LANCZOS199 except AttributeError:200 return Image.LANCZOS201 202# Make black transparent and scale opacity on a PIL Image203def make_black_transparent_and_adjust_opacity(204 img: Image.Image,205 opacity: float = 0.5,206 black_threshold: int = 10207) -> Image.Image:208 """209 - Near-black pixels (R,G,B < black_threshold) -> fully transparent.210 - All other pixels -> recolored to blue (0, 0, 255) with alpha scaled by `opacity`.211 Returns a PIL RGBA image.212 """213 img = img.convert("RGBA")214 pixels = img.getdata()215 216 new_data = []217 clamp = lambda v: max(0, min(255, v))218 219 for (r, g, b, a) in pixels:220 if r < black_threshold and g < black_threshold and b < black_threshold:221 # make black fully transparent222 new_data.append((0, 0, 0, 0))223 else:224 # recolor to blue and scale opacity225 new_alpha = clamp(int(a * opacity))226 new_data.append((0, 0, 255, new_alpha)) # pure blue227 228 img.putdata(new_data)229 return img230 231def overlay_images_rgba(232 base_img: Image.Image,233 overlay_img: Image.Image,234 resize_to_match: bool = True235) -> Image.Image:236 """237 Alpha-composites overlay_img onto base_img. Both treated as RGBA.238 If sizes differ and resize_to_match=True, overlay is resized to base size.239 """240 resample = _ensure_lanczos()241 base_rgba = base_img.convert("RGBA")242 over_rgba = overlay_img.convert("RGBA")243 244 if resize_to_match and base_rgba.size != over_rgba.size:245 over_rgba = over_rgba.resize(base_rgba.size, resample)246 247 # Alpha composite: overlay on top of base248 result = Image.alpha_composite(base_rgba, over_rgba)249 return result250 251 252# --- Fill Contour Function ---253def fill_contour_internal(image: Image.Image, contour_points) -> Image.Image:254 """255 Creates a binary mask where the contour is filled with white and everything else is black.256 257 Args:258 image (Image.Image): PIL Image used to get dimensions259 contour_points: Can be either:260 - List of [x, y] coordinates: [[100, 150], [200, 150], ...]261 - Comma-separated string: "373,295,372,296,373,297..."262 263 Returns:264 Image.Image: PIL Image (grayscale) with white-filled contour on black background265 266 Raises:267 ValueError: If contour_points is invalid or has fewer than 3 points268 """269 # If contour_points is a string, parse it270 if isinstance(contour_points, str):271 # Split by comma and convert to integers272 coords = [int(x.strip()) for x in contour_points.split(',')]273 274 # Check if we have an even number of coordinates275 if len(coords) % 2 != 0:276 raise ValueError("String contours must have even number of values (x,y pairs)")277 278 # Convert flat list to list of [x, y] pairs279 contour_points = [[coords[i], coords[i+1]] for i in range(0, len(coords), 2)]280 281 # Validate contour points282 if not isinstance(contour_points, list) or len(contour_points) < 3:283 raise ValueError("Contours must have at least 3 points")284 285 # Get image dimensions286 img_array = np.array(image)287 if len(img_array.shape) == 3:288 height, width = img_array.shape[:2]289 else:290 height, width = img_array.shape291 292 # Create a black mask293 mask = np.zeros((height, width), dtype=np.uint8)294 295 # Convert contour points to the correct shape for OpenCV296 contour = np.array(contour_points, dtype=np.int32).reshape((-1, 1, 2))297 298 # Fill the contour with white (255)299 cv2.fillPoly(mask, [contour], 255)300 301 # Convert to PIL Image302 return Image.fromarray(mask)303 304# route for Fill Contour 305 306@app.route('/fill_contour', methods=['POST'])307def fill_contour():308 """309 Fills a contour with white and makes everything else black.310 311 Body JSON:312 {313 "image": "<base64 or data URL>", # required, used to get dimensions314 "contours": [[x0,y0], [x1,y1], ...], # required, list of [x,y] points315 }316 317 Returns:318 {319 "success": true,320 "image": "<base64_png>", # binary mask (white contour, black background)321 "width": W,322 "height": H323 }324 """325 try:326 data = request.get_json()327 328 # Validate inputs329 if not data or 'image' not in data or 'contours' not in data:330 return jsonify({331 "success": False, 332 "error": "Provide both 'image' and 'contours'"333 }), 400334 335 # Decode the image336 pil_image = _decode_base64_to_pil(data['image'])337 338 # Get contour points339 contour_points = data['contours']340 341 # Generate the mask using the internal function342 mask_pil = fill_contour_internal(pil_image, contour_points)343 344 # Convert mask to base64345 mask_b64 = _pil_to_base64_png(mask_pil)346 347 width, height = mask_pil.size348 349 return jsonify({350 "success": True,351 "width": width,352 "height": height,353 "image": mask_b64 # base64 PNG (white contour on black background)354 })355 356 except ValueError as ve:357 return jsonify({358 "success": False,359 "error": str(ve)360 }), 400361 362 except Exception as e:363 return jsonify({364 "success": False,365 "error": str(e)366 }), 500367 368# --- New route: overlay two images with black→transparent on image2 ---369@app.route('/out_overlap_image', methods=['POST'])370def out_overlap_image():371 """372 Body JSON:373 {374 "image1": "<base64 or data URL>", # required, bottom layer375 "image2": "<base64 or data URL>", # required, top layer376 "opacity": 0.5, # optional, default 0.5 (applied to image2 non-black pixels)377 "black_threshold": 10, # optional, default 10 (RGB < threshold => transparent)378 "resize_to_match": true # optional, default true (resize image2 to image1 size)379 }380 Returns:381 { "success": true, "image": "<base64_png>", "width": W, "height": H }382 """383 try:384 data = request.get_json()385 if not data or 'image1' not in data or 'image2' not in data:386 return jsonify({"success": False, "error": "Provide image1 and image2"}), 400387 388 opacity = float(data.get('opacity', 0.5))389 black_threshold = int(data.get('black_threshold', 10))390 resize_to_match = bool(data.get('resize_to_match', True))391 392 # Decode inputs393 img1 = _decode_base64_to_pil(data['image1'])394 img2 = _decode_base64_to_pil(data['image2'])395 396 # Process image2: make black transparent + adjust opacity397 img2_processed = make_black_transparent_and_adjust_opacity(398 img2, opacity=opacity, black_threshold=black_threshold399 )400 401 # Overlay image2 over image1402 out = overlay_images_rgba(img1, img2_processed, resize_to_match=resize_to_match)403 404 out_b64 = _pil_to_base64_png(out)405 w, h = out.size406 return jsonify({407 "success": True,408 "width": w,409 "height": h,410 "image": out_b64 # base64 PNG (no data URL prefix)411 })412 413 except Exception as e:414 return jsonify({"success": False, "error": str(e)}), 500415 416 417 418@app.route('/predict', methods=['POST'])419def predict():420 try:421 data = request.get_json()422 423 if 'image' not in data:424 return jsonify({"error": "No image provided"}), 400425 426 # Get optional parameters427 hex_color = data.get('hex_color', '#0000FF') # Default blue428 tolerance = data.get('tolerance', 40) # Default tolerance429 430 # Decode base64 image431 image_data = data['image']432 if image_data.startswith('data:image'):433 image_data = image_data.split(',')[1]434 435 img_bytes = base64.b64decode(image_data)436 image = Image.open(io.BytesIO(img_bytes))437 438 # Detect blobs with color filtering439 result = detect_blobs_internal(image, hex_color, tolerance)440 return jsonify(result)441 442 except Exception as e:443 return jsonify({"error": str(e)}), 500444 445 446# --- /crop endpoint ---447@app.route('/crop', methods=['POST'])448def crop():449 try:450 data = request.get_json()451 if not data or 'image' not in data:452 return jsonify({"error": "No image provided"}), 400453 454 # Optional: allow custom output size (defaults to 1024)455 output_size = int(data.get('size', 1024))456 457 # Decode base64 (supports data URLs)458 image_data = data['image']459 if isinstance(image_data, str) and image_data.startswith('data:image'):460 image_data = image_data.split(',', 1)[1]461 462 img_bytes = base64.b64decode(image_data)463 pil_image = Image.open(io.BytesIO(img_bytes))464 465 result_image = center_crop_to_square(pil_image, output_size=output_size)466 467 # Encode back to base64 PNG468 buf = io.BytesIO()469 result_image.save(buf, format='PNG')470 buf.seek(0)471 out_b64 = base64.b64encode(buf.read()).decode('utf-8')472 473 return jsonify({474 "success": True,475 "width": output_size,476 "height": output_size,477 "image": out_b64, # base64 PNG without data URL prefix478 })479 except Exception as e:480 return jsonify({"success": False, "error": str(e)}), 500481 482 483 484# --- /sam2 endpoint ---485 486SPACE_URL = "https://augmentus-robotics-segment-anything-model-2.hf.space"487 488# Get HuggingFace token from environment variable489HF_TOKEN = os.getenv("HF_TOKEN") 490 491# Create client with authentication492from gradio_client import Client493client = Client(SPACE_URL, hf_token=HF_TOKEN)494 495# Optional: Add authentication middleware for your API endpoint496def require_api_key(f):497 @wraps(f)498 def decorated_function(*args, **kwargs):499 api_key = request.headers.get('X-API-Key')500 expected_key = os.getenv("API_KEY") # Your custom API key501 502 if expected_key and api_key != expected_key:503 return jsonify({"success": False, "error": "Invalid API key"}), 401504 return f(*args, **kwargs)505 return decorated_function506 507@app.post("/sam2pred")508# @require_api_key # Uncomment if you want to protect your endpoint509def sam2pred():510 try:511 data = request.get_json(force=True) or {}512 513 image_b64 = data["image_base64"]514 mask_b64 = data["mask_base64"]515 model = data.get("model_checkpoint", "large")516 alpha = float(data.get("alpha", 0.4))517 overlay_color = data.get("overlay_color", "0,0,255")518 return_metadata = bool(data.get("return_metadata", True))519 520 # Call the Space via gradio_client (joins the queue for you)521 resp = client.predict(522 image_b64,523 mask_b64,524 model,525 alpha,526 overlay_color,527 return_metadata,528 api_name="/predict",529 )530 531 # Your Space returns a JSON STRING; parse then re-emit clean JSON532 result = json.loads(resp) if isinstance(resp, str) else resp533 return jsonify(result)534 535 except KeyError as e:536 return jsonify({"success": False, "error": f"Missing field: {e.args[0]}"}), 400537 except Exception as e:538 return jsonify({"success": False, "error": str(e)}), 500539 540 541 542# --- / segment endpoint ---543@app.route('/segment', methods=['POST'])544def GeminiSegment():545 try:546 data = request.get_json()547 if not data or 'image' not in data or 'prompt' not in data:548 return jsonify({"error": "Missing 'image' or 'prompt' in request"}), 400549 550 # Extract request data551 image_data = data['image']552 object_to_segment = data['prompt']553 554 # Decode base64 image555 img_bytes = base64.b64decode(image_data)556 pil_image = PILImage.open(io.BytesIO(img_bytes)).convert("RGB")557 558 # Build dynamic prompt559 prompt = (560 f"Apply a pixel-perfect segmentation mask for {object_to_segment} in the image. "561 f"Output a JSON list of segmentation masks where each entry contains "562 f"the 2D bounding box in key 'box_2d', the segmentation mask in key 'mask', "563 f"and the text label in key 'label'."564 )565 566 # Run segmentation model567 binary_mask, segmentation_data = segment_image(prompt, pil_image)568 569 # Ensure mask is a valid image570 if isinstance(binary_mask, np.ndarray):571 if binary_mask.dtype != np.uint8:572 binary_mask = (binary_mask * 255).astype(np.uint8)573 if len(binary_mask.shape) == 2:574 binary_mask = cv2.cvtColor(binary_mask, cv2.COLOR_GRAY2RGB)575 else:576 return jsonify({"error": "Invalid mask output"}), 500577 578 # Encode mask as PNG base64579 _, buffer = cv2.imencode('.png', binary_mask)580 mask_base64 = base64.b64encode(buffer).decode('utf-8')581 582 return jsonify({583 "success": True,584 "mask_image": mask_base64585 })586 587 except Exception as e:588 return jsonify({"success": False, "error": str(e)}), 500589 590 591 592 593@app.route('/', methods=['GET'])594def home():595 return jsonify({596 "message": "Color-Based Blob Detection API",597 "endpoint": "/predict",598 "method": "POST",599 "parameters": {600 "image": "base64_encoded_image_string (required)",601 "hex_color": "hex color code, e.g. '#0000FF' (optional, default: #0000FF)",602 "tolerance": "color tolerance 0-255 (optional, default: 40)"603 },604 "example": {605 "image": "base64_string_here",606 "hex_color": "#0000FF",607 "tolerance": 40608 }609 })610 611@app.route('/health', methods=['GET'])612def health():613 return jsonify({"status": "healthy"})614 615if __name__ == '__main__':616 app.run(host='0.0.0.0', port=7860)