Files
FishServer/FishMeasure/utils/head_tail_alignment.py
2026-04-08 19:32:23 +08:00

114 lines
3.8 KiB
Python
Executable File

import numpy as np
def ensure_head_in_positive_x(points, bins=60):
"""
Ensure the fish head points toward +X and tail toward -X.
Algorithm:
1. Cut the fish into 4 parts along the X-axis.
2. For the first part (positive X side) and last part (negative X side),
divide into `bins` segments along Y-axis.
3. For each Y bin, compute thickness as |ymax - ymin|.
4. Calculate average thickness for the first and last parts.
5. The thinner part is the tail.
6. If tail lies at positive X, rotate 180° about Y-axis.
Args:
points (np.ndarray): Nx3 array of point coordinates.
bins (int): Number of bins along Y-axis for each part.
Returns:
tuple: (updated_points, flipped_flag)
"""
if points is None or len(points) == 0:
return points, False
pts = np.asarray(points)
x_vals = pts[:, 0]
y_vals = pts[:, 1]
x_min, x_max = x_vals.min(), x_vals.max()
if np.isclose(x_max, x_min):
# Degenerate along X; nothing to do
return pts, False
# Split fish into 4 parts along X-axis
x_range = x_max - x_min
x_quarter = x_range / 4.0
# First part: x >= x_max - x_quarter (positive X side, frontmost part)
first_part_mask = x_vals >= (x_max - x_quarter)
# Last part: x <= x_min + x_quarter (negative X side, rearmost part)
last_part_mask = x_vals <= (x_min + x_quarter)
if not np.any(first_part_mask) or not np.any(last_part_mask):
# Can't split properly
return pts, False
bins = max(10, int(bins)) # ensure reasonable bin count
def calculate_average_thickness(mask):
"""Calculate average thickness for points in this part.
For each part:
1. Get all points in this part
2. Divide Y-axis into bins
3. For each Y bin, compute thickness as |ymax - ymin| within that bin
4. Return average thickness across all Y bins
"""
pts_half = pts[mask]
if len(pts_half) == 0:
return float('inf')
y_half = pts_half[:, 1]
y_min_half, y_max_half = y_half.min(), y_half.max()
y_range = y_max_half - y_min_half
# Divide into bins along Y-axis
if np.isclose(y_range, 0):
return 0.0
bin_edges = np.linspace(y_min_half, y_max_half, bins + 1)
thicknesses = []
for i in range(bins):
left_edge = bin_edges[i]
right_edge = bin_edges[i + 1]
# Include right edge on final bin
if i == bins - 1:
y_mask = (y_half >= left_edge) & (y_half <= right_edge)
else:
y_mask = (y_half >= left_edge) & (y_half < right_edge)
if not np.any(y_mask):
continue
# For this Y bin, compute thickness as |ymax - ymin| within the bin
y_bin_values = y_half[y_mask]
thickness = float(np.abs(y_bin_values.max() - y_bin_values.min()))
thicknesses.append(thickness)
if not thicknesses:
return float('inf')
# Average thickness across all Y bins
return np.mean(thicknesses)
# Calculate average thickness for first and last parts
first_part_thickness = calculate_average_thickness(first_part_mask)
last_part_thickness = calculate_average_thickness(last_part_mask)
# The thinner part is the tail
tail_on_positive_x = first_part_thickness < last_part_thickness
if tail_on_positive_x:
# Tail is on +X; rotate 180° around Y-axis so head faces +X.
rotated = pts.copy()
rotated[:, 0] = -rotated[:, 0]
rotated[:, 2] = -rotated[:, 2]
return rotated, True
return pts, False