migration and splitting AareLC
This commit is contained in:
@@ -0,0 +1,12 @@
|
||||
"""Core functionality for AareLC ML Studio."""
|
||||
from .zmq_client import ZMQStreamClient
|
||||
from .inference_client import InferenceClient
|
||||
from .db_client import DatabaseClient
|
||||
from .image_processor import ImageProcessor
|
||||
|
||||
__all__ = [
|
||||
'ZMQStreamClient',
|
||||
'InferenceClient',
|
||||
'DatabaseClient',
|
||||
'ImageProcessor'
|
||||
]
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Computer vision logic utilities for edge detection and polygon refinement.
|
||||
|
||||
This module is maintained for backward compatibility.
|
||||
For new code, please use src.core.image_processor.ImageProcessor instead.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_refined_polygon(self, roi, p):
|
||||
"""
|
||||
Robust contour detection optimized for homogeneous backgrounds.
|
||||
Uses Adaptive Thresholding + Hole Filling.
|
||||
"""
|
||||
# 1. Preprocessing
|
||||
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Median blur is superior for preserving edges on homogeneous backgrounds
|
||||
blurred = cv2.medianBlur(gray, 5)
|
||||
|
||||
# 2. Adaptive Thresholding
|
||||
# block_size must be odd and > 1. We use morph_kernel slider to control this.
|
||||
block_size = max(3, p["morph"] if p["morph"] % 2 != 0 else p["morph"] + 1)
|
||||
|
||||
# Gaussian adaptive thresholding is very resilient to homogeneous background noise
|
||||
thresh = cv2.adaptiveThreshold(
|
||||
blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
||||
cv2.THRESH_BINARY_INV, block_size, 2
|
||||
)
|
||||
|
||||
# 3. Clean up noise
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
||||
opened = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
|
||||
|
||||
# 4. Fill Holes (Crucial for loops)
|
||||
# We find all contours and draw them filled on a mask
|
||||
cnts, _ = cv2.findContours(opened, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
|
||||
mask = np.zeros_like(opened)
|
||||
for c in cnts:
|
||||
cv2.drawContours(mask, [c], -1, 255, -1) # -1 thickness = fill
|
||||
|
||||
# 5. Final Contour Extraction
|
||||
# Now we take the largest filled blob
|
||||
final_cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
if not final_cnts:
|
||||
return np.array([]), opened
|
||||
|
||||
# Select largest blob by area
|
||||
cnt = max(final_cnts, key=cv2.contourArea)
|
||||
|
||||
# 6. Smooth the polygon based on the user's precision slider
|
||||
epsilon = p["eps"] * cv2.arcLength(cnt, True)
|
||||
approx = cv2.approxPolyDP(cnt, epsilon, True)
|
||||
|
||||
return approx, mask
|
||||
|
||||
|
||||
def measure_focus(raw_frame):
|
||||
"""
|
||||
Measure focus quality using Laplacian variance.
|
||||
|
||||
DEPRECATED: Use ImageProcessor.measure_focus() instead.
|
||||
|
||||
Args:
|
||||
raw_frame: Grayscale image
|
||||
|
||||
Returns:
|
||||
Focus score (higher = better focus)
|
||||
"""
|
||||
return cv2.Laplacian(raw_frame, cv2.CV_64F).var()
|
||||
@@ -0,0 +1,296 @@
|
||||
"""Client for database server communication."""
|
||||
import requests
|
||||
import numpy as np
|
||||
import cv2
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class DatabaseClient:
|
||||
"""Handles communication with database server for image retrieval."""
|
||||
|
||||
def __init__(self, server_url, shared_password="", download_base=None):
|
||||
"""
|
||||
Initialize database client.
|
||||
|
||||
Args:
|
||||
server_url: Base URL of database server (e.g., "http://localhost:8000")
|
||||
shared_password: Password for shared auth
|
||||
download_base: Base URL for image downloads (e.g., "http://localhost:8002/backend")
|
||||
"""
|
||||
self.server_url = server_url
|
||||
self.shared_password = shared_password
|
||||
self.download_base = download_base or "http://localhost:8002/backend"
|
||||
|
||||
def _get_headers(self):
|
||||
return {"x-shared-password": self.shared_password}
|
||||
|
||||
def fetch_next_raw_image(self, timeout=5):
|
||||
# 1. Get metadata from Dispatcher (via Tunnel 8001 or Direct 443)
|
||||
# Assuming you've set self.server_url to the dispatcher base
|
||||
meta_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/raw_next"
|
||||
r = requests.get(meta_url, headers=self._get_headers(), verify=False, timeout=timeout)
|
||||
|
||||
if r.status_code != 200:
|
||||
raise ValueError(f"No more images available (status {r.status_code})")
|
||||
|
||||
data = r.json()
|
||||
image_obj = data.get('image', {})
|
||||
|
||||
# 2. Construct the correct download URL
|
||||
filepath = image_obj.get('filepath')
|
||||
if not filepath:
|
||||
raise ValueError("Metadata received but no filepath found.")
|
||||
|
||||
# Use the configurable download_base
|
||||
download_url = f"{self.download_base.rstrip('/')}/{filepath.lstrip('/')}"
|
||||
|
||||
print(f"DEBUG: Fetching pixels from: {download_url}")
|
||||
img_res = requests.get(download_url, timeout=timeout, verify=False)
|
||||
|
||||
if img_res.status_code != 200:
|
||||
# Fallback debug
|
||||
raise ValueError(f"Could not retrieve image from {download_url} (status {img_res.status_code})")
|
||||
|
||||
# 3. Process the content
|
||||
content = img_res.content
|
||||
|
||||
# Use OpenCV to decode
|
||||
nparr = np.frombuffer(content, np.uint8)
|
||||
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||||
|
||||
if img is None:
|
||||
# If standard decode fails, check if it's a RAW Bayer file (which your system uses)
|
||||
# A 4MP Bayer file is exactly 4194304 bytes or 4177936 bytes
|
||||
if len(content) in [4194304, 4177936]:
|
||||
try:
|
||||
# Attempt raw Bayer recovery (2048x2048 or 2048x2040)
|
||||
h = 2048
|
||||
w = len(content) // h
|
||||
raw = np.frombuffer(content, np.uint8).reshape((h, w))
|
||||
img = cv2.cvtColor(raw, cv2.COLOR_BAYER_RG2BGR)
|
||||
print(f"DEBUG: Successfully decoded as raw Bayer {w}x{h}")
|
||||
except:
|
||||
pass
|
||||
|
||||
if img is None:
|
||||
raise ValueError(f"Failed to decode image data. size={len(content)} bytes. Header: {content[:4].hex()}")
|
||||
|
||||
# Return both image and metadata
|
||||
return img, data
|
||||
|
||||
def set_server_url(self, url):
|
||||
"""Update server URL."""
|
||||
self.server_url = url
|
||||
|
||||
def get_server_url(self):
|
||||
"""Get current server URL."""
|
||||
return self.server_url
|
||||
|
||||
def set_shared_password(self, password):
|
||||
"""Update the shared password."""
|
||||
self.shared_password = password
|
||||
|
||||
def set_download_base(self, url):
|
||||
"""Update the image download base URL."""
|
||||
self.download_base = url
|
||||
|
||||
def save_annotation(self, image_id, segments, username=None, timeout=10):
|
||||
"""
|
||||
Save YOLO segmentation annotations to database.
|
||||
|
||||
Args:
|
||||
image_id: ID of the image in the database
|
||||
segments: List of dicts with 'class_id' and 'points' (normalized YOLO format)
|
||||
username: Optional username to attribute the annotation to
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
Response dict with status and annotation_id
|
||||
"""
|
||||
save_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/{image_id}"
|
||||
|
||||
payload = {
|
||||
"segments": segments,
|
||||
"username": username
|
||||
}
|
||||
|
||||
print(f"DEBUG: Saving annotation to: {save_url}")
|
||||
print(f"DEBUG: Payload: {len(segments)} segments")
|
||||
|
||||
response = requests.post(
|
||||
save_url,
|
||||
json=payload,
|
||||
headers=self._get_headers(),
|
||||
verify=False,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"Failed to save annotation: {response.status_code} - {response.text}")
|
||||
|
||||
return response.json()
|
||||
|
||||
def skip_image(self, image_id, timeout=5):
|
||||
"""
|
||||
Mark an image as skipped in the database.
|
||||
|
||||
Args:
|
||||
image_id: ID of the image to skip
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
Response dict with status
|
||||
"""
|
||||
skip_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/{image_id}/skip"
|
||||
|
||||
print(f"DEBUG: Skipping image: {skip_url}")
|
||||
|
||||
response = requests.post(
|
||||
skip_url,
|
||||
headers=self._get_headers(),
|
||||
verify=False,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"Failed to skip image: {response.status_code} - {response.text}")
|
||||
|
||||
return response.json()
|
||||
|
||||
|
||||
def sync_local_annotations(
|
||||
self,
|
||||
since=None,
|
||||
until=None,
|
||||
include_unannotated=False,
|
||||
dry_run=False,
|
||||
limit=5000,
|
||||
timeout=120,
|
||||
):
|
||||
"""Retrieve DB sync items and write missing files locally on this machine."""
|
||||
sync_url = f"{self.server_url.rstrip('/')}/dispatcher/protected_router/annotations/sync/local"
|
||||
|
||||
params = {
|
||||
"include_unannotated": str(bool(include_unannotated)).lower(),
|
||||
"limit": int(limit),
|
||||
}
|
||||
if since:
|
||||
params["since"] = since
|
||||
if until:
|
||||
params["until"] = until
|
||||
|
||||
print(f"DEBUG: Syncing local files from: {sync_url}")
|
||||
print(f"DEBUG: Sync params: {params}")
|
||||
|
||||
response = requests.post(
|
||||
sync_url,
|
||||
params=params,
|
||||
headers=self._get_headers(),
|
||||
verify=False,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise ValueError(
|
||||
f"Failed to sync local files: {response.status_code} - {response.text}"
|
||||
)
|
||||
data = response.json()
|
||||
items = data.get("items", [])
|
||||
|
||||
images_dir = Path("data") / "img_sets" / "seg_training" / "images"
|
||||
labels_dir = Path("data") / "img_sets" / "seg_training" / "labels"
|
||||
if not dry_run:
|
||||
images_dir.mkdir(parents=True, exist_ok=True)
|
||||
labels_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
stats = {
|
||||
"scanned": len(items),
|
||||
"image_copied": 0,
|
||||
"image_existing": 0,
|
||||
"image_failed": 0,
|
||||
"image_skipped_no_label": 0,
|
||||
"label_copied": 0,
|
||||
"label_existing": 0,
|
||||
"label_failed": 0,
|
||||
"dry_run": bool(dry_run),
|
||||
"since": since,
|
||||
"until": until,
|
||||
}
|
||||
failures = []
|
||||
|
||||
for item in items:
|
||||
image_id = item.get("image_id")
|
||||
image_url = item.get("image_url")
|
||||
segments = item.get("segments") or []
|
||||
|
||||
if image_id is None:
|
||||
stats["image_failed"] += 1
|
||||
failures.append("missing image_id in sync item")
|
||||
continue
|
||||
|
||||
image_path = images_dir / f"img_{int(image_id)}.jpg"
|
||||
label_path = labels_dir / f"img_{int(image_id)}.txt"
|
||||
|
||||
lines = []
|
||||
for seg in segments:
|
||||
points = seg.get("points")
|
||||
if (
|
||||
not isinstance(points, list)
|
||||
or len(points) < 6
|
||||
or len(points) % 2 != 0
|
||||
):
|
||||
continue
|
||||
class_id = int(seg.get("class_id", 0))
|
||||
lines.append(f"{class_id} " + " ".join(f"{float(p):.6f}" for p in points))
|
||||
|
||||
label_ready = False
|
||||
if label_path.exists():
|
||||
stats["label_existing"] += 1
|
||||
label_ready = True
|
||||
elif not lines:
|
||||
stats["label_failed"] += 1
|
||||
stats["image_skipped_no_label"] += 1
|
||||
failures.append(
|
||||
f"image_id={image_id}: no valid segmentation points for label"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
try:
|
||||
if not dry_run:
|
||||
label_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
stats["label_copied"] += 1
|
||||
label_ready = True
|
||||
except Exception as exc:
|
||||
stats["label_failed"] += 1
|
||||
failures.append(
|
||||
f"image_id={image_id}: failed writing {label_path} ({exc})"
|
||||
)
|
||||
continue
|
||||
|
||||
if not label_ready:
|
||||
stats["image_skipped_no_label"] += 1
|
||||
continue
|
||||
|
||||
if image_path.exists():
|
||||
stats["image_existing"] += 1
|
||||
else:
|
||||
try:
|
||||
img_res = requests.get(
|
||||
image_url,
|
||||
timeout=min(timeout, 30),
|
||||
verify=False,
|
||||
)
|
||||
if img_res.status_code != 200 or not img_res.content:
|
||||
raise ValueError(
|
||||
f"image download failed status={img_res.status_code}"
|
||||
)
|
||||
if not dry_run:
|
||||
image_path.write_bytes(img_res.content)
|
||||
stats["image_copied"] += 1
|
||||
except Exception as exc:
|
||||
stats["image_failed"] += 1
|
||||
failures.append(
|
||||
f"image_id={image_id}: failed downloading {image_url} ({exc})"
|
||||
)
|
||||
|
||||
return {"status": "ok", "stats": stats, "failures": failures[:100]}
|
||||
@@ -0,0 +1,586 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class ImageProcessor:
|
||||
def __init__(self):
|
||||
self.edge_map = None
|
||||
self.edge_map_shape = None
|
||||
self.roi_offset = (0, 0) # Store (x1, y1) of the box
|
||||
|
||||
def compute_edge_map(self, image, box=None):
|
||||
"""
|
||||
Compute edge map, optionally restricted to a specific box.
|
||||
"""
|
||||
if image is None:
|
||||
return None
|
||||
|
||||
# 1. Handle ROI crop if box is provided
|
||||
if box:
|
||||
x1, y1, x2, y2 = int(box['x1']), int(box['y1']), int(box['x2']), int(box['y2'])
|
||||
# Add small padding to capture edges right at the boundary
|
||||
h, w = image.shape[:2]
|
||||
p = 10
|
||||
x1, y1 = max(0, x1-p), max(0, y1-p)
|
||||
x2, y2 = min(w, x2+p), min(h, y2+p)
|
||||
|
||||
working_img = image[y1:y2, x1:x2]
|
||||
self.roi_offset = (x1, y1)
|
||||
else:
|
||||
working_img = image
|
||||
self.roi_offset = (0, 0)
|
||||
|
||||
# 2. Convert and process local ROI
|
||||
gray = cv2.cvtColor(working_img, cv2.COLOR_BGR2GRAY)
|
||||
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
||||
|
||||
adaptive = cv2.adaptiveThreshold(
|
||||
blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
||||
cv2.THRESH_BINARY_INV, 11, 2
|
||||
)
|
||||
canny = cv2.Canny(blurred, 30, 100)
|
||||
edge_map = cv2.bitwise_or(adaptive, canny)
|
||||
|
||||
kernel = np.ones((3,3), np.uint8)
|
||||
self.edge_map = cv2.dilate(edge_map, kernel, iterations=1)
|
||||
self.edge_map_shape = working_img.shape[:2]
|
||||
return self.edge_map
|
||||
|
||||
def find_nearest_edge(self, x, y, radius=20):
|
||||
"""
|
||||
Find nearest edge within the stored ROI.
|
||||
"""
|
||||
if self.edge_map is None:
|
||||
return x, y
|
||||
|
||||
# Translate global mouse coords to ROI local coords
|
||||
lx, ly = int(x - self.roi_offset[0]), int(y - self.roi_offset[1])
|
||||
h, w = self.edge_map.shape
|
||||
|
||||
# Bounds check
|
||||
x0, x1 = max(0, lx - radius), min(w - 1, lx + radius)
|
||||
y0, y1 = max(0, ly - radius), min(h - 1, ly + radius)
|
||||
|
||||
region = self.edge_map[y0:y1+1, x0:x1+1]
|
||||
edge_points = np.argwhere(region > 128)
|
||||
|
||||
if len(edge_points) == 0:
|
||||
return x, y
|
||||
|
||||
# Find closest point in local ROI space
|
||||
edge_points_local = edge_points + [y0, x0]
|
||||
distances = np.linalg.norm(edge_points_local - [ly, lx], axis=1)
|
||||
closest_idx = np.argmin(distances)
|
||||
closest_y, closest_x = edge_points_local[closest_idx]
|
||||
|
||||
# Translate back to global image coords
|
||||
return int(closest_x + self.roi_offset[0]), int(closest_y + self.roi_offset[1])
|
||||
|
||||
@staticmethod
|
||||
def get_refined_polygon_watershed(roi, p, manual_mask=None):
|
||||
"""
|
||||
Watershed-based segmentation using manual strokes as markers.
|
||||
More intuitive than GrabCut: background strokes act as barriers.
|
||||
"""
|
||||
h, w = roi.shape[:2]
|
||||
print(f" Watershed: roi_shape={roi.shape}")
|
||||
if h < 10 or w < 10:
|
||||
print(f" Watershed: ROI too small ({h}x{w})")
|
||||
return np.array([]), np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
# Convert to grayscale for processing
|
||||
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Compute gradient magnitude (edges)
|
||||
sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
|
||||
sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
|
||||
gradient = np.sqrt(sobelx**2 + sobely**2)
|
||||
gradient = cv2.normalize(gradient, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
|
||||
|
||||
# Create markers for watershed
|
||||
markers = np.zeros((h, w), dtype=np.int32)
|
||||
|
||||
if manual_mask is not None:
|
||||
# Match sizes
|
||||
if manual_mask.shape[:2] != (h, w):
|
||||
print(f" Watershed: Resizing manual_mask from {manual_mask.shape[:2]} to ({h}, {w})")
|
||||
manual_mask = cv2.resize(manual_mask, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
has_bg = np.any(manual_mask == 0)
|
||||
has_fg = np.any(manual_mask == 1)
|
||||
print(f" Watershed: has_bg={has_bg}, has_fg={has_fg}")
|
||||
|
||||
if has_bg or has_fg:
|
||||
# Background strokes = marker 1 (will be excluded)
|
||||
if has_bg:
|
||||
markers[manual_mask == 0] = 1
|
||||
|
||||
# Foreground strokes = marker 2 (will be included)
|
||||
if has_fg:
|
||||
fg_pixels = np.sum(manual_mask == 1)
|
||||
print(f" Watershed: Using {fg_pixels} foreground pixels")
|
||||
markers[manual_mask == 1] = 2
|
||||
|
||||
# If only foreground painted, automatically add background border
|
||||
if not has_bg:
|
||||
print(f" Watershed: No background painted - using border as background hint")
|
||||
border_width = max(3, min(w, h) // 30)
|
||||
markers[:border_width, :] = 1
|
||||
markers[-border_width:, :] = 1
|
||||
markers[:, :border_width] = 1
|
||||
markers[:, -border_width:] = 1
|
||||
else:
|
||||
# No foreground marked - use center region
|
||||
center_x, center_y = w // 2, h // 2
|
||||
radius = min(w, h) // 6
|
||||
y_indices, x_indices = np.ogrid[:h, :w]
|
||||
center_mask = (x_indices - center_x)**2 + (y_indices - center_y)**2 <= radius**2
|
||||
markers[center_mask] = 2
|
||||
print(f" Watershed: Using center circle as foreground hint")
|
||||
|
||||
# Apply watershed
|
||||
print(f" Watershed: Running watershed algorithm...")
|
||||
markers = cv2.watershed(roi, markers)
|
||||
print(f" Watershed: Watershed complete, unique markers: {np.unique(markers)}")
|
||||
|
||||
# Extract foreground (marker == 2)
|
||||
binary_mask = np.where(markers == 2, 255, 0).astype(np.uint8)
|
||||
fg_area = np.sum(binary_mask > 0)
|
||||
print(f" Watershed: Binary mask has {fg_area} foreground pixels")
|
||||
else:
|
||||
# No manual strokes - use automatic thresholding
|
||||
binary_mask = ImageProcessor._get_auto_mask(roi, p)
|
||||
else:
|
||||
# No manual mask - use automatic thresholding
|
||||
binary_mask = ImageProcessor._get_auto_mask(roi, p)
|
||||
|
||||
# Morphological cleanup
|
||||
k = max(1, p.get("morph", 7))
|
||||
k = k if k % 2 != 0 else k + 1
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||||
binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
|
||||
binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel)
|
||||
|
||||
# Find contours
|
||||
contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
print(f" Watershed: Found {len(contours)} contours")
|
||||
if not contours:
|
||||
print(f" Watershed: No contours found!")
|
||||
return np.array([]), binary_mask
|
||||
|
||||
# Get largest contour
|
||||
cnt = max(contours, key=cv2.contourArea)
|
||||
area = cv2.contourArea(cnt)
|
||||
print(f" Watershed: Largest contour area: {area}")
|
||||
|
||||
# Smooth polygon
|
||||
epsilon = p.get("eps", 0.01) * cv2.arcLength(cnt, True)
|
||||
approx = cv2.approxPolyDP(cnt, epsilon, True)
|
||||
print(f" Watershed: Polygon has {len(approx)} points")
|
||||
|
||||
return approx, binary_mask
|
||||
|
||||
@staticmethod
|
||||
def _get_auto_mask(roi, p):
|
||||
"""Automatic mask generation without manual strokes."""
|
||||
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
||||
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
||||
|
||||
# Adaptive thresholding
|
||||
adaptive = cv2.adaptiveThreshold(
|
||||
blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
||||
cv2.THRESH_BINARY_INV, 11, 2
|
||||
)
|
||||
|
||||
# Canny edges
|
||||
canny = cv2.Canny(blurred, 30, p.get("high", 150))
|
||||
|
||||
# Combine
|
||||
binary_mask = cv2.bitwise_or(adaptive, canny)
|
||||
|
||||
return binary_mask
|
||||
|
||||
@staticmethod
|
||||
def get_refined_polygon(roi, p, manual_mask=None):
|
||||
"""
|
||||
Refinement using Interactive GrabCut.
|
||||
Optimized for homogeneous backgrounds by using the ROI border as a background seed.
|
||||
"""
|
||||
h, w = roi.shape[:2]
|
||||
if h < 10 or w < 10: return np.array([]), np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
pad_w, pad_h = int(w * 0.05), int(h * 0.05)
|
||||
rect = (pad_w, pad_h, w - 2 * pad_w, h - 2 * pad_h)
|
||||
|
||||
# Initialize mask with proper GrabCut values
|
||||
if manual_mask is not None:
|
||||
# Match sizes to prevent IndexErrors
|
||||
if manual_mask.shape[:2] != roi.shape[:2]:
|
||||
manual_mask = cv2.resize(manual_mask, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# Check what strokes we have
|
||||
has_bg = np.any(manual_mask == 0)
|
||||
has_fg = np.any(manual_mask == 1)
|
||||
|
||||
# GrabCut needs BOTH foreground and background samples
|
||||
if has_bg and not has_fg:
|
||||
# User only painted background - use rectangle center as foreground hint
|
||||
print(" GrabCut mode: Background only - adding center as foreground hint")
|
||||
mask = np.full((h, w), cv2.GC_PR_BGD, dtype=np.uint8) # Start with probable background
|
||||
mask[manual_mask == 0] = cv2.GC_BGD # Hard background from user
|
||||
# Set center region as probable foreground
|
||||
center_x, center_y = w // 2, h // 2
|
||||
radius = min(w, h) // 4
|
||||
y_indices, x_indices = np.ogrid[:h, :w]
|
||||
center_mask = (x_indices - center_x)**2 + (y_indices - center_y)**2 <= radius**2
|
||||
mask[center_mask] = cv2.GC_PR_FGD
|
||||
print(f" BG pixels: {np.sum(mask == cv2.GC_BGD)}, FG hint pixels: {np.sum(mask == cv2.GC_PR_FGD)}")
|
||||
elif has_fg and not has_bg:
|
||||
# User only painted foreground - use border as background hint
|
||||
mask = np.full((h, w), cv2.GC_PR_FGD, dtype=np.uint8) # Start with probable foreground
|
||||
mask[manual_mask == 1] = cv2.GC_FGD # Hard foreground from user
|
||||
# Set border as probable background
|
||||
border_width = max(2, min(w, h) // 20)
|
||||
mask[:border_width, :] = cv2.GC_PR_BGD
|
||||
mask[-border_width:, :] = cv2.GC_PR_BGD
|
||||
mask[:, :border_width] = cv2.GC_PR_BGD
|
||||
mask[:, -border_width:] = cv2.GC_PR_BGD
|
||||
elif has_bg and has_fg:
|
||||
# User painted both - use their constraints directly
|
||||
mask = np.full((h, w), cv2.GC_PR_FGD, dtype=np.uint8)
|
||||
mask[manual_mask == 0] = cv2.GC_BGD
|
||||
mask[manual_mask == 1] = cv2.GC_FGD
|
||||
else:
|
||||
# No strokes at all - fall back to rectangle mode
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
mode = cv2.GC_INIT_WITH_RECT
|
||||
iters = max(1, min(2, int(p.get("high", 150) / 75)))
|
||||
bgdModel = np.zeros((1, 65), np.float64)
|
||||
fgdModel = np.zeros((1, 65), np.float64)
|
||||
try:
|
||||
cv2.grabCut(roi, mask, rect, bgdModel, fgdModel, iters, mode)
|
||||
except Exception as e:
|
||||
print(f"GrabCut failed: {e}")
|
||||
return np.array([]), np.zeros((h, w), dtype=np.uint8)
|
||||
binary_mask = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8') * 255
|
||||
k = max(1, p.get("morph", 7))
|
||||
k = k if k % 2 != 0 else k + 1
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||||
binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
|
||||
contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
if not contours:
|
||||
return np.array([]), binary_mask
|
||||
cnt = max(contours, key=cv2.contourArea)
|
||||
epsilon = p.get("eps", 0.01) * cv2.arcLength(cnt, True)
|
||||
approx = cv2.approxPolyDP(cnt, epsilon, True)
|
||||
return approx, binary_mask
|
||||
|
||||
mode = cv2.GC_INIT_WITH_MASK
|
||||
# More iterations when we have manual guidance
|
||||
iters = 5
|
||||
else:
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
mode = cv2.GC_INIT_WITH_RECT
|
||||
iters = max(1, min(2, int(p.get("high", 150) / 75)))
|
||||
|
||||
bgdModel = np.zeros((1, 65), np.float64)
|
||||
fgdModel = np.zeros((1, 65), np.float64)
|
||||
|
||||
try:
|
||||
print(f" Running GrabCut: mode={mode}, iters={iters}, roi_shape={roi.shape}")
|
||||
cv2.grabCut(roi, mask, rect, bgdModel, fgdModel, iters, mode)
|
||||
print(f" GrabCut succeeded")
|
||||
except Exception as e:
|
||||
print(f" GrabCut failed: {e}")
|
||||
return np.array([]), np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
# Convert to binary mask
|
||||
binary_mask = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8') * 255
|
||||
|
||||
# 4. Clean up the mask using the morphology slider
|
||||
k = max(1, p.get("morph", 7))
|
||||
k = k if k % 2 != 0 else k + 1
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||||
binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
|
||||
|
||||
# 5. Find the largest contour
|
||||
contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
if not contours:
|
||||
return np.array([]), binary_mask
|
||||
|
||||
cnt = max(contours, key=cv2.contourArea)
|
||||
|
||||
# 6. Smooth the polygon based on the epsilon slider
|
||||
epsilon = p.get("eps", 0.01) * cv2.arcLength(cnt, True)
|
||||
approx = cv2.approxPolyDP(cnt, epsilon, True)
|
||||
|
||||
return approx, binary_mask
|
||||
|
||||
def draw_detections(self, image, detections, class_params, class_colors, active_class=None,
|
||||
show_debug_edges=False, refine_mode=False, selected_idx=-1, selected_island_idx=0,
|
||||
selected_island_kind="outer", hidden_classes=None,
|
||||
segmentation_method='watershed', annotation_line_width=2,
|
||||
show_boxes=True, show_segments=True):
|
||||
"""
|
||||
Draw detection boxes and polygons.
|
||||
Supports manual editing by checking for existing 'poly' in detections.
|
||||
"""
|
||||
draw_img = image.copy()
|
||||
hidden_classes = hidden_classes or set()
|
||||
|
||||
# Show global edge map overlay if requested
|
||||
if show_debug_edges and self.edge_map is not None:
|
||||
# Create colored edge overlay
|
||||
edge_colored = cv2.cvtColor(self.edge_map, cv2.COLOR_GRAY2BGR)
|
||||
edge_colored[:, :, 0] = 0 # Remove blue channel
|
||||
edge_colored[:, :, 1] = 0 # Remove green channel
|
||||
# Keep only red channel for edges
|
||||
draw_img = cv2.addWeighted(draw_img, 0.7, edge_colored, 0.3, 0)
|
||||
|
||||
for i, det in enumerate(detections):
|
||||
is_selected = (i == selected_idx)
|
||||
try:
|
||||
box = det.get('box', det)
|
||||
x1, y1, x2, y2 = int(box.get('x1', 0)), int(box.get('y1', 0)), int(box.get('x2', 0)), int(
|
||||
box.get('y2', 0))
|
||||
label = det.get('name', det.get('label', 'unknown'))
|
||||
except:
|
||||
continue
|
||||
|
||||
# Skip hidden classes
|
||||
if label in hidden_classes:
|
||||
continue
|
||||
|
||||
color = class_colors.get(label, (255, 0, 255))
|
||||
params = class_params.get(label) or next(iter(class_params.values()), {})
|
||||
base_thickness = max(1, int(annotation_line_width))
|
||||
thickness = base_thickness + 1 if is_selected else base_thickness
|
||||
|
||||
# 1. Draw Bounding Box
|
||||
if show_boxes:
|
||||
cv2.rectangle(draw_img, (x1, y1), (x2, y2), color, thickness)
|
||||
cv2.putText(draw_img, f"{label} {'(EDITING)' if is_selected else ''}",
|
||||
(x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
|
||||
|
||||
if is_selected and show_boxes:
|
||||
# DRAW RESIZE HANDLES (Small white squares at all 4 corners)
|
||||
handle_size = max(3, base_thickness + 2)
|
||||
corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
|
||||
for (cx, cy) in corners:
|
||||
cv2.rectangle(draw_img, (cx - handle_size, cy - handle_size),
|
||||
(cx + handle_size, cy + handle_size), (255, 255, 255), -1)
|
||||
cv2.rectangle(draw_img, (cx - handle_size, cy - handle_size),
|
||||
(cx + handle_size, cy + handle_size), color, max(1, base_thickness - 1))
|
||||
|
||||
# 2. Handle Polygon (Manual or Refined)
|
||||
polygons = []
|
||||
if isinstance(det.get('polygons'), list):
|
||||
for poly_item in det.get('polygons'):
|
||||
arr = np.array(poly_item, dtype=np.float32).reshape(-1, 1, 2)
|
||||
if len(arr) >= 3:
|
||||
polygons.append(arr)
|
||||
poly = det.get('poly')
|
||||
if not polygons and poly is not None and len(poly) >= 3:
|
||||
polygons = [np.array(poly, dtype=np.float32).reshape(-1, 1, 2)]
|
||||
holes = []
|
||||
if isinstance(det.get('holes'), list):
|
||||
for hole_item in det.get('holes'):
|
||||
arr = np.array(hole_item, dtype=np.float32).reshape(-1, 1, 2)
|
||||
if len(arr) >= 3:
|
||||
holes.append(arr)
|
||||
manual_mask = det.get('manual_mask') # Retrieve user brush strokes stored in the detection
|
||||
|
||||
# Regenerate polygon if: missing, empty, or needs refinement
|
||||
suppress_auto_polygon = bool(det.get('suppress_auto_polygon', False))
|
||||
should_regenerate = show_segments and refine_mode and y2 > y1 and x2 > x1 and (
|
||||
len(polygons) == 0
|
||||
) and not suppress_auto_polygon
|
||||
|
||||
if should_regenerate:
|
||||
print(f" Should regenerate: method={segmentation_method}, manual_mask={'present' if manual_mask is not None else 'absent'}")
|
||||
roi = image[y1:y2, x1:x2]
|
||||
if roi.size > 0:
|
||||
# Use selected segmentation method
|
||||
if segmentation_method == 'watershed':
|
||||
generated_poly, debug_mask = self.get_refined_polygon_watershed(roi, params, manual_mask)
|
||||
elif segmentation_method == 'grabcut':
|
||||
generated_poly, debug_mask = self.get_refined_polygon(roi, params, manual_mask)
|
||||
elif segmentation_method == 'smart':
|
||||
print(f" Calling smart ridge-snapping segmentation...")
|
||||
generated_poly, debug_mask = self.get_refined_polygon_livewire(roi, params, manual_mask)
|
||||
else:
|
||||
generated_poly, debug_mask = self.get_refined_polygon_watershed(roi, params, manual_mask)
|
||||
|
||||
if len(generated_poly) > 0:
|
||||
# CRITICAL FIX: Ensure the polygon is stored back in the detection
|
||||
generated_poly = generated_poly.astype(np.float32)
|
||||
det['poly'] = generated_poly
|
||||
det['polygons'] = [generated_poly]
|
||||
polygons = [generated_poly]
|
||||
print(f" Stored smart polygon with {len(generated_poly)} points in detection")
|
||||
|
||||
b, g, r = color
|
||||
colored_mask = cv2.merge([
|
||||
np.where(debug_mask > 0, b, 0).astype(np.uint8),
|
||||
np.where(debug_mask > 0, g, 0).astype(np.uint8),
|
||||
np.where(debug_mask > 0, r, 0).astype(np.uint8)
|
||||
])
|
||||
draw_img[y1:y2, x1:x2] = cv2.addWeighted(draw_img[y1:y2, x1:x2], 0.6, colored_mask, 0.4, 0)
|
||||
|
||||
# 3. Draw the Polygon and interactive points
|
||||
if show_segments and polygons:
|
||||
for poly_idx, poly_item in enumerate(polygons):
|
||||
pts = (poly_item.reshape((-1, 1, 2)) + [x1, y1]).astype(np.int32)
|
||||
cv2.polylines(draw_img, [pts], isClosed=True, color=color, thickness=thickness)
|
||||
|
||||
if is_selected:
|
||||
vertex_fill_radius = max(2, base_thickness + 1)
|
||||
vertex_border_radius = vertex_fill_radius + 1
|
||||
active_idx = max(0, min(int(selected_island_idx), len(polygons) - 1))
|
||||
if selected_island_kind == "outer" and poly_idx == active_idx:
|
||||
for p_idx, p in enumerate(pts):
|
||||
cv2.circle(draw_img, tuple(p[0]), vertex_fill_radius, (255, 255, 255), -1)
|
||||
cv2.circle(draw_img, tuple(p[0]), vertex_border_radius, color, max(1, base_thickness - 1))
|
||||
cv2.putText(
|
||||
draw_img,
|
||||
f"outer {active_idx + 1}/{len(polygons)}",
|
||||
(x1, y2 + 16),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.45,
|
||||
color,
|
||||
1
|
||||
)
|
||||
|
||||
if show_segments and holes:
|
||||
hole_color = (0, 165, 255)
|
||||
for hole_idx, hole_item in enumerate(holes):
|
||||
pts = (hole_item.reshape((-1, 1, 2)) + [x1, y1]).astype(np.int32)
|
||||
cv2.polylines(draw_img, [pts], isClosed=True, color=hole_color, thickness=max(1, thickness - 1))
|
||||
if is_selected:
|
||||
vertex_fill_radius = max(2, base_thickness + 1)
|
||||
vertex_border_radius = vertex_fill_radius + 1
|
||||
active_idx = max(0, min(int(selected_island_idx), len(holes) - 1))
|
||||
if selected_island_kind == "hole" and hole_idx == active_idx:
|
||||
for p in pts:
|
||||
cv2.circle(draw_img, tuple(p[0]), vertex_fill_radius, (255, 255, 255), -1)
|
||||
cv2.circle(draw_img, tuple(p[0]), vertex_border_radius, hole_color, max(1, base_thickness - 1))
|
||||
cv2.putText(
|
||||
draw_img,
|
||||
f"hole {active_idx + 1}/{len(holes)}",
|
||||
(x1, y2 + 32),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.45,
|
||||
hole_color,
|
||||
1
|
||||
)
|
||||
|
||||
# 4. Visualize Brush Strokes with colored overlay
|
||||
if show_segments and is_selected and manual_mask is not None:
|
||||
# Get the actual ROI region from the image
|
||||
roi_region = draw_img[y1:y2, x1:x2]
|
||||
roi_h, roi_w = roi_region.shape[:2]
|
||||
|
||||
# Ensure manual_mask matches ROI size
|
||||
if manual_mask.shape[:2] != (roi_h, roi_w):
|
||||
manual_mask_resized = cv2.resize(manual_mask, (roi_w, roi_h), interpolation=cv2.INTER_NEAREST)
|
||||
else:
|
||||
manual_mask_resized = manual_mask
|
||||
|
||||
# Create colored overlay for brush strokes
|
||||
overlay = np.zeros((roi_h, roi_w, 3), dtype=np.uint8)
|
||||
|
||||
# Background strokes in RED (semi-transparent)
|
||||
bg_mask = manual_mask_resized == 0
|
||||
overlay[bg_mask] = [0, 0, 255] # Red in BGR
|
||||
|
||||
# Foreground strokes in GREEN (semi-transparent)
|
||||
fg_mask = manual_mask_resized == 1
|
||||
overlay[fg_mask] = [0, 255, 0] # Green in BGR
|
||||
|
||||
# Only blend where we have brush strokes (not the "probable" areas)
|
||||
mask_combined = (bg_mask | fg_mask).astype(np.uint8)
|
||||
if mask_combined.sum() > 0:
|
||||
# Create alpha channel effect
|
||||
alpha = 0.4 # 40% opacity for the overlay
|
||||
blended = cv2.addWeighted(roi_region, 1 - alpha, overlay, alpha, 0)
|
||||
# Only apply blending where strokes exist
|
||||
roi_region[mask_combined > 0] = blended[mask_combined > 0]
|
||||
draw_img[y1:y2, x1:x2] = roi_region
|
||||
|
||||
return draw_img
|
||||
|
||||
@staticmethod
|
||||
def measure_focus(raw_frame):
|
||||
return cv2.Laplacian(raw_frame, cv2.CV_64F).var()
|
||||
|
||||
@staticmethod
|
||||
def draw_fps_overlay(image, fps_in, fps_out, fps_pred=0.0):
|
||||
"""Draw input/output stream FPS overlay on an annotated frame."""
|
||||
if image is None:
|
||||
return image
|
||||
text = f"FPS in (pserv): {fps_in:.1f} | out (unified): {fps_out:.1f} | pred: {fps_pred:.1f}"
|
||||
cv2.putText(
|
||||
image,
|
||||
text,
|
||||
(10, 88),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.6,
|
||||
(0, 255, 255),
|
||||
2
|
||||
)
|
||||
return image
|
||||
|
||||
@staticmethod
|
||||
def get_refined_polygon_livewire(roi, p, manual_mask=None):
|
||||
"""
|
||||
Smart selection using Gradient ridge-snapping.
|
||||
"""
|
||||
h, w = roi.shape[:2]
|
||||
if h < 10 or w < 10: return np.array([]), np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
# 1. Edge Detection
|
||||
gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
|
||||
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
||||
|
||||
# High sensitivity Canny
|
||||
low_t = max(10, p.get("high", 150) // 3)
|
||||
edges = cv2.Canny(blurred, low_t, p.get("high", 150))
|
||||
|
||||
# 2. Use Manual Mask to block areas
|
||||
if manual_mask is not None:
|
||||
if manual_mask.shape[:2] != (h, w):
|
||||
manual_mask = cv2.resize(manual_mask, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
# Remove edges in background-painted areas
|
||||
edges[manual_mask == 0] = 0
|
||||
|
||||
# 3. Create a 'snapping' surface using distance transform
|
||||
# This creates a 'valley' at every edge
|
||||
inv_edges = cv2.bitwise_not(edges)
|
||||
dist_trans = cv2.distanceTransform(inv_edges, cv2.DIST_L2, 5)
|
||||
|
||||
# 4. Generate a base mask (Watershed is actually a great seed for this)
|
||||
_, seed_mask = ImageProcessor.get_refined_polygon_watershed(roi, p, manual_mask)
|
||||
|
||||
# 5. Shrink/Expand seed to snap to the distance transform ridge
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
||||
snapped_mask = cv2.morphologyEx(seed_mask, cv2.MORPH_CLOSE, kernel)
|
||||
|
||||
# 6. Find final contours on the snapped mask
|
||||
# Use CHAIN_APPROX_SIMPLE for reliable point extraction
|
||||
contours, _ = cv2.findContours(snapped_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
if not contours:
|
||||
print(" Smart: No contours found in snapped_mask")
|
||||
return np.array([]), snapped_mask
|
||||
|
||||
cnt = max(contours, key=cv2.contourArea)
|
||||
|
||||
# 7. Apply smoothing and ensure float32 format for storage
|
||||
eps_factor = max(0.0001, p.get("eps", 0.005))
|
||||
epsilon = eps_factor * cv2.arcLength(cnt, True)
|
||||
approx = cv2.approxPolyDP(cnt, epsilon, True)
|
||||
|
||||
# Ensure the polygon is a float32 array in (N, 1, 2) shape
|
||||
poly_out = approx.astype(np.float32)
|
||||
print(f" Smart: Generated polygon with {len(poly_out)} points")
|
||||
|
||||
return poly_out, snapped_mask
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Client for ML inference server communication."""
|
||||
import requests
|
||||
import cv2
|
||||
import json
|
||||
import numpy as np
|
||||
|
||||
|
||||
class InferenceClient:
|
||||
"""Handles communication with ML inference server."""
|
||||
|
||||
def __init__(self, server_url):
|
||||
"""
|
||||
Initialize inference client.
|
||||
|
||||
Args:
|
||||
server_url: Base URL of inference server (e.g., "http://mx-ml.psi.ch:8002")
|
||||
"""
|
||||
self.server_url = server_url
|
||||
|
||||
def predict(self, image, timeout=10):
|
||||
"""
|
||||
Send image to server for prediction.
|
||||
|
||||
Args:
|
||||
image: OpenCV image (BGR format)
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
List of detection results, or None on error
|
||||
|
||||
Raises:
|
||||
requests.RequestException: On connection errors
|
||||
ValueError: On invalid response
|
||||
"""
|
||||
_, buf = cv2.imencode('.jpg', image)
|
||||
url = f"{self.server_url}/predict/"
|
||||
files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
|
||||
|
||||
response = requests.post(url, files=files, timeout=timeout)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
return data.get('results', [])
|
||||
else:
|
||||
raise ValueError(f"Server returned status {response.status_code}")
|
||||
|
||||
def segment_anything(self, image, x, y, timeout=15):
|
||||
"""
|
||||
Request a high-precision mask from the SAM backend.
|
||||
"""
|
||||
_, buf = cv2.imencode('.jpg', image)
|
||||
# Use trailing slash to match the server's /sam/ route exactly
|
||||
url = f"{self.server_url.rstrip('/')}/sam/"
|
||||
params = {"point_x": int(x), "point_y": int(y)}
|
||||
files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
|
||||
|
||||
response = requests.post(url, files=files, params=params, timeout=timeout)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
raise ValueError(f"SAM error: {response.status_code} - {response.text}")
|
||||
|
||||
def segment_anything_box(self, image, x1, y1, x2, y2, timeout=15):
|
||||
"""
|
||||
Request a high-precision mask from the SAM backend using a bounding box prompt.
|
||||
"""
|
||||
_, buf = cv2.imencode('.jpg', image)
|
||||
url = f"{self.server_url.rstrip('/')}/sam_box/"
|
||||
params = {
|
||||
"x1": int(x1),
|
||||
"y1": int(y1),
|
||||
"x2": int(x2),
|
||||
"y2": int(y2)
|
||||
}
|
||||
files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
|
||||
|
||||
response = requests.post(url, files=files, params=params, timeout=timeout)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
raise ValueError(f"SAM Box error: {response.status_code} - {response.text}")
|
||||
|
||||
def segment_multi_box(self, image, boxes, timeout=30):
|
||||
"""
|
||||
Request masks for multiple boxes in one call.
|
||||
boxes: List of [x1, y1, x2, y2]
|
||||
"""
|
||||
_, buf = cv2.imencode('.jpg', image)
|
||||
url = f"{self.server_url.rstrip('/')}/sam_multi_box/"
|
||||
params = {"boxes_json": json.dumps(boxes)}
|
||||
files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
|
||||
|
||||
response = requests.post(url, files=files, params=params, timeout=timeout)
|
||||
if response.status_code == 200:
|
||||
return response.json().get('results', [])
|
||||
else:
|
||||
raise ValueError(f"SAM Multi Box error: {response.status_code}")
|
||||
|
||||
def predict_depth(self, image, timeout=20, colorize=True, depth_model=None, depth_max_side=0):
|
||||
"""
|
||||
Request depth map for an image.
|
||||
|
||||
Returns:
|
||||
OpenCV image array decoded from PNG response.
|
||||
"""
|
||||
_, buf = cv2.imencode('.jpg', image)
|
||||
url = f"{self.server_url.rstrip('/')}/depth/"
|
||||
params = {"colorize": "true" if colorize else "false"}
|
||||
if depth_model:
|
||||
params["depth_model"] = str(depth_model)
|
||||
if isinstance(depth_max_side, int) and depth_max_side > 0:
|
||||
params["depth_max_side"] = str(depth_max_side)
|
||||
files = {'file': ('image.jpg', buf.tobytes(), 'image/jpeg')}
|
||||
|
||||
response = requests.post(url, files=files, params=params, timeout=timeout)
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"Depth error: {response.status_code} - {response.text}")
|
||||
if "image/png" not in response.headers.get("content-type", ""):
|
||||
try:
|
||||
err = response.json().get("error", response.text)
|
||||
except Exception:
|
||||
err = response.text
|
||||
raise ValueError(f"Depth error: {err}")
|
||||
|
||||
depth_png = cv2.imdecode(np.frombuffer(response.content, np.uint8), cv2.IMREAD_UNCHANGED)
|
||||
if depth_png is None:
|
||||
raise ValueError("Depth response decode failed.")
|
||||
return depth_png
|
||||
|
||||
def send_training_data(self, image, metadata, timeout=10):
|
||||
"""
|
||||
Send annotated image to training pipeline.
|
||||
|
||||
Args:
|
||||
image: OpenCV image (BGR format)
|
||||
metadata: Dict containing params and detections
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
Response object
|
||||
|
||||
Raises:
|
||||
requests.RequestException: On connection errors
|
||||
"""
|
||||
_, buf = cv2.imencode('.jpg', image)
|
||||
url = f"{self.server_url}/train/"
|
||||
files = {'file': ('img.jpg', buf.tobytes(), 'image/jpeg')}
|
||||
data = {"metadata": json.dumps(metadata)}
|
||||
|
||||
response = requests.post(url, files=files, data=data, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
|
||||
def set_server_url(self, url):
|
||||
"""Update server URL."""
|
||||
self.server_url = url
|
||||
|
||||
def get_server_url(self):
|
||||
"""Get current server URL."""
|
||||
return self.server_url
|
||||
@@ -0,0 +1,201 @@
|
||||
"""ZMQ client for receiving live image streams with detection results."""
|
||||
import zmq
|
||||
import json
|
||||
import numpy as np
|
||||
import cv2
|
||||
import threading
|
||||
|
||||
|
||||
class ZMQStreamClient:
|
||||
"""Handles ZMQ subscription for live detector stream."""
|
||||
|
||||
def __init__(self, zmq_url, on_frame_callback, on_focus_callback=None, use_cuda=True):
|
||||
"""
|
||||
Initialize ZMQ stream client.
|
||||
|
||||
Args:
|
||||
zmq_url: ZMQ server URL (e.g., "tcp://localhost:9091")
|
||||
on_frame_callback: Callback function(frame, detections, metadata) called on new frame
|
||||
on_focus_callback: Optional callback function(focus_score) for focus updates
|
||||
use_cuda: Use OpenCV CUDA path when available on this client
|
||||
"""
|
||||
self.zmq_url = zmq_url
|
||||
self.on_frame_callback = on_frame_callback
|
||||
self.on_focus_callback = on_focus_callback
|
||||
self.focus_enabled = on_focus_callback is not None
|
||||
self.use_cuda = bool(use_cuda)
|
||||
self.cuda_enabled = False
|
||||
self._cuda_lap_filter = None
|
||||
self.streaming = False
|
||||
self.worker_thread = None
|
||||
|
||||
if self.use_cuda:
|
||||
try:
|
||||
if hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0:
|
||||
# Create reusable Laplacian filter for focus score on GPU.
|
||||
self._cuda_lap_filter = cv2.cuda.createLaplacianFilter(cv2.CV_8UC1, cv2.CV_32F, 3)
|
||||
self.cuda_enabled = True
|
||||
except Exception:
|
||||
self.cuda_enabled = False
|
||||
self._cuda_lap_filter = None
|
||||
self._cuda_fallback_logged = False
|
||||
|
||||
def start(self):
|
||||
"""Start streaming in background thread."""
|
||||
if self.streaming:
|
||||
return
|
||||
self.streaming = True
|
||||
self.worker_thread = threading.Thread(target=self._worker, daemon=True)
|
||||
self.worker_thread.start()
|
||||
|
||||
def stop(self):
|
||||
"""Stop streaming."""
|
||||
self.streaming = False
|
||||
if self.worker_thread:
|
||||
self.worker_thread.join(timeout=1.0)
|
||||
|
||||
def is_streaming(self):
|
||||
"""Check if currently streaming."""
|
||||
return self.streaming
|
||||
|
||||
def enable_focus(self):
|
||||
"""Enable focus-score computation and callbacks."""
|
||||
self.focus_enabled = True
|
||||
|
||||
def disable_focus(self):
|
||||
"""Disable focus-score computation and callbacks."""
|
||||
self.focus_enabled = False
|
||||
|
||||
def _worker(self):
|
||||
"""Background worker that listens to ZMQ stream."""
|
||||
print(f"DEBUG: Starting Unified ZMQ worker ({self.zmq_url}) cuda={self.cuda_enabled}", flush=True)
|
||||
if self.use_cuda:
|
||||
if self.cuda_enabled:
|
||||
print("INFO: OpenCV CUDA path enabled for focus scoring", flush=True)
|
||||
else:
|
||||
print("INFO: OpenCV CUDA requested but unavailable, using CPU path", flush=True)
|
||||
ctx = zmq.Context()
|
||||
sub = ctx.socket(zmq.SUB)
|
||||
sub.connect(self.zmq_url)
|
||||
sub.setsockopt_string(zmq.SUBSCRIBE, "")
|
||||
sub.setsockopt(zmq.RCVBUF, 128 * 1024 * 1024)
|
||||
poller = zmq.Poller()
|
||||
poller.register(sub, zmq.POLLIN)
|
||||
|
||||
while self.streaming:
|
||||
try:
|
||||
socks = dict(poller.poll(10))
|
||||
if sub in socks:
|
||||
# Drain queue to get latest message
|
||||
latest_msg = None
|
||||
while True:
|
||||
try:
|
||||
latest_msg = sub.recv_multipart(zmq.NOBLOCK)
|
||||
except zmq.Again:
|
||||
break
|
||||
|
||||
if latest_msg:
|
||||
self._process_message(latest_msg)
|
||||
except Exception as e:
|
||||
print(f"ZMQ Error: {e}")
|
||||
break
|
||||
|
||||
sub.close()
|
||||
ctx.term()
|
||||
|
||||
def _process_message(self, msg_parts):
|
||||
"""Process received ZMQ multipart message."""
|
||||
try:
|
||||
# Last part is image data, rest might contain header
|
||||
img_data = msg_parts[-1]
|
||||
header = None
|
||||
|
||||
# Try to find JSON header in message parts
|
||||
for part in msg_parts[:-1]:
|
||||
try:
|
||||
if not part:
|
||||
continue
|
||||
header = json.loads(part.decode('utf-8'))
|
||||
if isinstance(header, dict) and 'shape' in header:
|
||||
break
|
||||
header = None
|
||||
except:
|
||||
continue
|
||||
|
||||
if not header:
|
||||
return
|
||||
|
||||
# Extract detections
|
||||
detections = header.get('boxes', [])
|
||||
if isinstance(detections, list):
|
||||
for det in detections:
|
||||
if not isinstance(det, dict):
|
||||
continue
|
||||
poly = det.get('poly')
|
||||
if poly is None:
|
||||
continue
|
||||
try:
|
||||
poly_arr = np.asarray(poly, dtype=np.float32)
|
||||
if poly_arr.ndim == 2 and poly_arr.shape[1] == 2 and poly_arr.shape[0] >= 3:
|
||||
det['poly'] = poly_arr.reshape(-1, 1, 2)
|
||||
else:
|
||||
det.pop('poly', None)
|
||||
except Exception:
|
||||
det.pop('poly', None)
|
||||
|
||||
# Decode image payload (raw BGR or Bayer from unified stream server).
|
||||
pixel_format = str(header.get('pixel_format', 'BGR')).upper()
|
||||
shape = header.get('shape', [])
|
||||
if pixel_format == 'BAYER':
|
||||
if len(shape) != 2:
|
||||
return
|
||||
h, w = int(shape[0]), int(shape[1])
|
||||
raw = np.frombuffer(img_data, np.uint8).reshape((h, w))
|
||||
bayer_pattern = str(header.get('bayer_pattern', 'RGGB')).upper()
|
||||
demosaic_map = {
|
||||
'RGGB': cv2.COLOR_BAYER_RG2BGR,
|
||||
'GBRG': cv2.COLOR_BAYER_GB2BGR,
|
||||
'GRBG': cv2.COLOR_BAYER_GR2BGR,
|
||||
'BGGR': cv2.COLOR_BAYER_BG2BGR,
|
||||
}
|
||||
demosaic_code = demosaic_map.get(bayer_pattern)
|
||||
if demosaic_code is None:
|
||||
return
|
||||
img = cv2.cvtColor(raw, demosaic_code)
|
||||
else:
|
||||
if len(shape) != 3:
|
||||
return
|
||||
h, w, c = int(shape[0]), int(shape[1]), int(shape[2])
|
||||
if c != 3:
|
||||
return
|
||||
img = np.frombuffer(img_data, np.uint8).reshape((h, w, c))
|
||||
|
||||
# Calculate focus score
|
||||
if self.focus_enabled and self.on_focus_callback:
|
||||
if self.cuda_enabled:
|
||||
try:
|
||||
gpu_img = cv2.cuda_GpuMat()
|
||||
gpu_img.upload(img)
|
||||
gpu_gray = cv2.cuda.cvtColor(gpu_img, cv2.COLOR_BGR2GRAY)
|
||||
gpu_lap = self._cuda_lap_filter.apply(gpu_gray)
|
||||
lap = gpu_lap.download()
|
||||
focus_score = float(lap.var())
|
||||
except Exception:
|
||||
if not self._cuda_fallback_logged:
|
||||
print("WARN: CUDA focus path failed at runtime, falling back to CPU", flush=True)
|
||||
self._cuda_fallback_logged = True
|
||||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
focus_score = cv2.Laplacian(gray, cv2.CV_64F).var()
|
||||
else:
|
||||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
focus_score = cv2.Laplacian(gray, cv2.CV_64F).var()
|
||||
self.on_focus_callback(focus_score)
|
||||
|
||||
# Call frame callback. Keep backward compatibility with 2-arg callbacks.
|
||||
try:
|
||||
self.on_frame_callback(img, detections, header)
|
||||
except TypeError:
|
||||
self.on_frame_callback(img, detections)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing ZMQ message: {e}")
|
||||
Reference in New Issue
Block a user