472 lines
17 KiB
Python
472 lines
17 KiB
Python
import json
|
|
import cv2
|
|
import numpy as np
|
|
import zmq
|
|
import time
|
|
import threading
|
|
from typing import Dict, Any
|
|
|
|
|
|
class ImageStatsReceiver:
|
|
def __init__(self, zmq_url: str):
|
|
"""Initialize ZeroMQ receiver for image statistics calculation"""
|
|
context = zmq.Context()
|
|
self.__socket = context.socket(zmq.SUB)
|
|
self.__socket.setsockopt(zmq.SUBSCRIBE, b"")
|
|
self.__socket.setsockopt(zmq.RCVTIMEO, 1000) # 1 second timeout
|
|
self.__socket.setsockopt(zmq.LINGER, 0) # Don't linger on close
|
|
|
|
try:
|
|
self.__socket.connect(zmq_url)
|
|
print(f"Connected to ZeroMQ socket: {zmq_url}")
|
|
except Exception as e:
|
|
print(f"Failed to connect to {zmq_url}: {e}")
|
|
raise
|
|
|
|
self.running = False
|
|
self.latest_stats = None
|
|
self.stats_lock = threading.Lock()
|
|
self.message_count = 0
|
|
self.connection_attempts = 0
|
|
self.last_message_time = None
|
|
self.connection_attempts = 0
|
|
self.last_message_time = None
|
|
|
|
def calculate_projections(self, image: np.ndarray) -> Dict[str, Any]:
|
|
"""Calculate X and Y projections of the image"""
|
|
# Convert to grayscale if color image
|
|
if len(image.shape) == 3:
|
|
gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
|
else:
|
|
gray_image = image.copy()
|
|
|
|
# Calculate projections (sum along each axis)
|
|
x_projection = np.mean(gray_image, axis=0) # Sum along Y axis -> X profile
|
|
y_projection = np.mean(gray_image, axis=1) # Sum along X axis -> Y profile
|
|
|
|
# Also calculate max projections (useful for finding peaks)
|
|
x_projection_max = np.max(gray_image, axis=0)
|
|
y_projection_max = np.max(gray_image, axis=1)
|
|
|
|
# Find peak positions and values
|
|
x_peak_idx = np.argmax(x_projection)
|
|
y_peak_idx = np.argmax(y_projection)
|
|
x_peak_value = x_projection[x_peak_idx]
|
|
y_peak_value = y_projection[y_peak_idx]
|
|
|
|
# Calculate centroid positions (intensity-weighted center)
|
|
x_coords = np.arange(len(x_projection))
|
|
y_coords = np.arange(len(y_projection))
|
|
|
|
x_centroid = (
|
|
np.sum(x_coords * x_projection) / np.sum(x_projection)
|
|
if np.sum(x_projection) > 0
|
|
else len(x_projection) / 2
|
|
)
|
|
y_centroid = (
|
|
np.sum(y_coords * y_projection) / np.sum(y_projection)
|
|
if np.sum(y_projection) > 0
|
|
else len(y_projection) / 2
|
|
)
|
|
|
|
# Calculate FWHM (Full Width at Half Maximum) approximation
|
|
def calculate_fwhm(profile):
|
|
peak_val = np.max(profile)
|
|
half_max = peak_val / 2
|
|
indices = np.where(profile >= half_max)[0]
|
|
if len(indices) > 0:
|
|
return indices[-1] - indices[0]
|
|
return 0
|
|
|
|
x_fwhm = calculate_fwhm(x_projection)
|
|
y_fwhm = calculate_fwhm(y_projection)
|
|
|
|
return {
|
|
"x_projection": x_projection.tolist(),
|
|
"y_projection": y_projection.tolist(),
|
|
"x_projection_max": x_projection_max.tolist(),
|
|
"y_projection_max": y_projection_max.tolist(),
|
|
"x_peak_position": int(x_peak_idx),
|
|
"y_peak_position": int(y_peak_idx),
|
|
"x_peak_value": float(x_peak_value),
|
|
"y_peak_value": float(y_peak_value),
|
|
"x_centroid": float(x_centroid),
|
|
"y_centroid": float(y_centroid),
|
|
"x_fwhm": float(x_fwhm),
|
|
"y_fwhm": float(y_fwhm),
|
|
"x_coords": list(range(len(x_projection))),
|
|
"y_coords": list(range(len(y_projection))),
|
|
}
|
|
|
|
def calculate_radial_integration(self, image: np.ndarray, num_bins: int = 50) -> Dict[str, Any]:
|
|
"""Calculate radial integration of the image"""
|
|
# Convert to grayscale if color image
|
|
if len(image.shape) == 3:
|
|
gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
|
else:
|
|
gray_image = image.copy()
|
|
|
|
# Get image center
|
|
center_y, center_x = np.array(gray_image.shape) // 2
|
|
|
|
# Create coordinate grids
|
|
y_coords, x_coords = np.ogrid[: gray_image.shape[0], : gray_image.shape[1]]
|
|
|
|
# Calculate distance from center for each pixel
|
|
distances = np.sqrt((x_coords - center_x) ** 2 + (y_coords - center_y) ** 2)
|
|
|
|
# Maximum possible radius (to corner of image)
|
|
max_radius = np.sqrt((gray_image.shape[0] / 2) ** 2 + (gray_image.shape[1] / 2) ** 2)
|
|
|
|
# Create radial bins
|
|
r_bins = np.linspace(0, max_radius, num_bins + 1)
|
|
r_centers = (r_bins[:-1] + r_bins[1:]) / 2
|
|
|
|
# Calculate radial profile
|
|
radial_profile = []
|
|
radial_std = []
|
|
pixel_counts = []
|
|
|
|
for i in range(num_bins):
|
|
# Create mask for current radial bin
|
|
mask = (distances >= r_bins[i]) & (distances < r_bins[i + 1])
|
|
|
|
if np.any(mask):
|
|
# Get pixel values in this radial bin
|
|
pixel_values = gray_image[mask]
|
|
radial_profile.append(float(np.mean(pixel_values)))
|
|
radial_std.append(float(np.std(pixel_values)))
|
|
pixel_counts.append(int(np.sum(mask)))
|
|
else:
|
|
radial_profile.append(0.0)
|
|
radial_std.append(0.0)
|
|
pixel_counts.append(0)
|
|
|
|
return {
|
|
"r_centers": r_centers.tolist(),
|
|
"radial_profile": radial_profile,
|
|
"radial_std": radial_std,
|
|
"pixel_counts": pixel_counts,
|
|
"max_radius": float(max_radius),
|
|
"center": [int(center_x), int(center_y)],
|
|
"num_bins": num_bins,
|
|
}
|
|
|
|
def calculate_image_stats(self, image: np.ndarray) -> Dict[str, Any]:
|
|
"""Calculate comprehensive statistics for an image"""
|
|
image_float = image.astype(np.float64)
|
|
|
|
stats = {
|
|
"timestamp": time.time(),
|
|
"shape": image.shape,
|
|
"dtype": str(image.dtype),
|
|
"mean": float(np.mean(image_float)),
|
|
"std": float(np.std(image_float)),
|
|
"median": float(np.median(image_float)),
|
|
"min": float(np.min(image_float)),
|
|
"max": float(np.max(image_float)),
|
|
"message_count": self.message_count,
|
|
}
|
|
|
|
# Calculate per-channel stats if color image
|
|
if len(image.shape) == 3 and image.shape[2] > 1:
|
|
for channel in range(image.shape[2]):
|
|
channel_data = image_float[:, :, channel]
|
|
stats[f"mean_ch{channel}"] = float(np.mean(channel_data))
|
|
stats[f"std_ch{channel}"] = float(np.std(channel_data))
|
|
stats[f"median_ch{channel}"] = float(np.median(channel_data))
|
|
|
|
radial_data = self.calculate_radial_integration(image)
|
|
stats["radial"] = radial_data
|
|
|
|
# projection_data = self.calculate_projections(image)
|
|
# stats['projections'] = projection_data
|
|
|
|
return stats
|
|
|
|
def process_message(self):
|
|
"""Process a single ZeroMQ message"""
|
|
try:
|
|
r = self.__socket.recv_multipart(zmq.NOBLOCK)
|
|
if len(r) != 2:
|
|
return None
|
|
|
|
meta, data = r
|
|
header = json.loads(meta)
|
|
header_shape = header["shape"]
|
|
self.message_count += 1
|
|
|
|
if header["type"] == "uint8" and len(header_shape) == 2:
|
|
# Bayer image
|
|
bayer_image = np.frombuffer(data, dtype=np.uint8)
|
|
bayer_image = bayer_image.reshape(header_shape)
|
|
|
|
# Convert to RGB for main stats
|
|
rgb_image = cv2.cvtColor(bayer_image, cv2.COLOR_BAYER_GB2RGB)
|
|
rgb_image = rgb_image[:, ::-1, :].copy()
|
|
|
|
stats = self.calculate_image_stats(rgb_image)
|
|
stats["image_type"] = "rgb_from_bayer"
|
|
|
|
with self.stats_lock:
|
|
self.latest_stats = stats
|
|
|
|
return stats
|
|
|
|
elif header["type"] == "uint8" and len(header_shape) == 3:
|
|
# RGB image
|
|
rgb_image = np.frombuffer(data, dtype=np.uint8)
|
|
rgb_image = rgb_image.reshape(header_shape)
|
|
|
|
stats = self.calculate_image_stats(rgb_image)
|
|
stats["image_type"] = "rgb_direct"
|
|
|
|
with self.stats_lock:
|
|
self.latest_stats = stats
|
|
|
|
return stats
|
|
|
|
except zmq.Again:
|
|
return None
|
|
except Exception as e:
|
|
print(f"Error processing ZMQ message: {e}")
|
|
return None
|
|
|
|
def receiver_thread(self):
|
|
"""Background thread for receiving messages"""
|
|
while self.running:
|
|
self.process_message()
|
|
time.sleep(0.001) # Small sleep to prevent busy waiting
|
|
|
|
def print_stats_thread(self):
|
|
"""Background thread for printing stats every second"""
|
|
while self.running:
|
|
time.sleep(5.0) # Print once per second
|
|
|
|
with self.stats_lock:
|
|
if self.latest_stats:
|
|
self.print_formatted_stats(self.latest_stats)
|
|
else:
|
|
print(f"[{time.strftime('%H:%M:%S')}] Waiting for image data...")
|
|
|
|
def print_formatted_stats(self, stats: Dict[str, Any]):
|
|
"""Print formatted statistics in a compact terminal format"""
|
|
timestamp = time.strftime("%H:%M:%S", time.localtime(stats["timestamp"]))
|
|
|
|
# Compact one-line format
|
|
print(
|
|
f"[{timestamp}] #{stats['message_count']:4d} | "
|
|
f"Shape: {stats['shape']} | "
|
|
f"Mean: {stats['mean']:6.1f} | "
|
|
f"Std: {stats['std']:6.1f} | "
|
|
f"Median: {stats['median']:6.1f} | "
|
|
f"Range: [{stats['min']:3.0f}-{stats['max']:3.0f}]",
|
|
flush=True,
|
|
)
|
|
|
|
if "radial" in stats:
|
|
radial = stats["radial"]
|
|
center_intensity = radial["radial_profile"][0] if radial["radial_profile"] else 0
|
|
edge_intensity = radial["radial_profile"][-1] if radial["radial_profile"] else 0
|
|
peak_radius_idx = np.argmax(radial["radial_profile"]) if radial["radial_profile"] else 0
|
|
peak_radius = radial["r_centers"][peak_radius_idx] if radial["r_centers"] else 0
|
|
|
|
print(
|
|
f"{'':21} Radial: Center={center_intensity:.1f} | "
|
|
f"Edge={edge_intensity:.1f} | "
|
|
f"Peak@r={peak_radius:.1f} | "
|
|
f"Center=({radial['center'][0]},{radial['center'][1]})",
|
|
flush=True,
|
|
)
|
|
|
|
# Print projection summary
|
|
if "projections" in stats:
|
|
proj = stats["projections"]
|
|
print(
|
|
f"{'':21} X-Profile: Peak@{proj['x_peak_position']}({proj['x_peak_value']:.1f}) | "
|
|
f"Centroid={proj['x_centroid']:.1f} | FWHM={proj['x_fwhm']:.1f}",
|
|
flush=True,
|
|
)
|
|
print(
|
|
f"{'':21} Y-Profile: Peak@{proj['y_peak_position']}({proj['y_peak_value']:.1f}) | "
|
|
f"Centroid={proj['y_centroid']:.1f} | FWHM={proj['y_fwhm']:.1f}",
|
|
flush=True,
|
|
)
|
|
|
|
# Optional: Print per-channel stats if available
|
|
if any(k.startswith("mean_ch") for k in stats.keys()):
|
|
channels = []
|
|
i = 0
|
|
while f"mean_ch{i}" in stats:
|
|
channels.append(f"Ch{i}({stats[f'mean_ch{i}']:.1f})")
|
|
i += 1
|
|
if channels:
|
|
print(f"{'':21} Channels: {' '.join(channels)}", flush=True)
|
|
|
|
def get_radial_profile_summary(self):
|
|
"""Get a summary of the current radial profile"""
|
|
with self.stats_lock:
|
|
if self.latest_stats and "radial" in self.latest_stats:
|
|
radial = self.latest_stats["radial"]
|
|
return {
|
|
"r_centers": radial["r_centers"],
|
|
"radial_profile": radial["radial_profile"],
|
|
"center_intensity": radial["radial_profile"][0]
|
|
if radial["radial_profile"]
|
|
else 0,
|
|
"edge_intensity": radial["radial_profile"][-1]
|
|
if radial["radial_profile"]
|
|
else 0,
|
|
"max_intensity_radius": radial["r_centers"][np.argmax(radial["radial_profile"])]
|
|
if radial["radial_profile"]
|
|
else 0,
|
|
"max_intensity_value": max(radial["radial_profile"])
|
|
if radial["radial_profile"]
|
|
else 0,
|
|
}
|
|
return None
|
|
|
|
def save_radial_profile_to_file(self, filename: str = None):
|
|
"""Save the current radial profile to a file"""
|
|
if filename is None:
|
|
filename = f"radial_profile_{int(time.time())}.txt"
|
|
|
|
with self.stats_lock:
|
|
if self.latest_stats and "radial" in self.latest_stats:
|
|
radial = self.latest_stats["radial"]
|
|
|
|
with open(filename, "w") as f:
|
|
f.write("# Radial Integration Profile\n")
|
|
f.write(f"# Timestamp: {time.ctime(self.latest_stats['timestamp'])}\n")
|
|
f.write(f"# Image shape: {self.latest_stats['shape']}\n")
|
|
f.write(f"# Center: {radial['center']}\n")
|
|
f.write("# Radius(pixels)\tMean_Intensity\tStd_Intensity\tPixel_Count\n")
|
|
|
|
for i in range(len(radial["r_centers"])):
|
|
f.write(
|
|
f"{radial['r_centers'][i]:.2f}\t"
|
|
f"{radial['radial_profile'][i]:.2f}\t"
|
|
f"{radial['radial_std'][i]:.2f}\t"
|
|
f"{radial['pixel_counts'][i]}\n"
|
|
)
|
|
|
|
print(f"Radial profile saved to {filename}")
|
|
return filename
|
|
|
|
print("No radial profile data available")
|
|
return None
|
|
|
|
def start(self):
|
|
"""Start the background threads"""
|
|
self.running = True
|
|
|
|
# Start receiver thread
|
|
self.receiver_thread_obj = threading.Thread(target=self.receiver_thread, daemon=True)
|
|
self.receiver_thread_obj.start()
|
|
|
|
# Start stats printing thread
|
|
self.print_thread_obj = threading.Thread(target=self.print_stats_thread, daemon=True)
|
|
self.print_thread_obj.start()
|
|
|
|
print("Image stats receiver started - printing every 1 second")
|
|
|
|
def stop(self):
|
|
"""Stop the receiver and all threads"""
|
|
self.running = False
|
|
if self.__socket:
|
|
self.__socket.close()
|
|
print("Image stats receiver stopped")
|
|
|
|
|
|
# Global stats receiver instance
|
|
stats_receiver = None
|
|
|
|
|
|
def start_image_stats_receiver(zmq_url: str = "tcp://localhost:5555"):
|
|
"""Start the image stats receiver in background"""
|
|
global stats_receiver
|
|
|
|
if stats_receiver is None or not stats_receiver.running:
|
|
try:
|
|
stats_receiver = ImageStatsReceiver(zmq_url)
|
|
stats_receiver.start()
|
|
print(f"Started image statistics monitoring on {zmq_url}")
|
|
except Exception as e:
|
|
print(f"Failed to start image stats receiver: {e}")
|
|
stats_receiver = None
|
|
|
|
|
|
def stop_image_stats_receiver():
|
|
"""Stop the image stats receiver"""
|
|
global stats_receiver
|
|
|
|
if stats_receiver and stats_receiver.running:
|
|
stats_receiver.stop()
|
|
stats_receiver = None
|
|
|
|
|
|
def get_latest_image_stats():
|
|
"""Get the latest image statistics"""
|
|
global stats_receiver
|
|
|
|
if stats_receiver and stats_receiver.latest_stats:
|
|
with stats_receiver.stats_lock:
|
|
return stats_receiver.latest_stats.copy()
|
|
return None
|
|
|
|
|
|
def get_latest_image_stats():
|
|
"""Get the latest image statistics including radial profile"""
|
|
global stats_receiver
|
|
|
|
if stats_receiver and stats_receiver.latest_stats:
|
|
with stats_receiver.stats_lock:
|
|
return stats_receiver.latest_stats.copy()
|
|
return None
|
|
|
|
|
|
def get_radial_profile():
|
|
"""Get just the radial profile data"""
|
|
global stats_receiver
|
|
|
|
if stats_receiver:
|
|
return stats_receiver.get_radial_profile_summary()
|
|
return None
|
|
|
|
|
|
def save_current_radial_profile(filename: str = None):
|
|
"""Save the current radial profile to a file"""
|
|
global stats_receiver
|
|
|
|
if stats_receiver:
|
|
return stats_receiver.save_radial_profile_to_file(filename)
|
|
return None
|
|
|
|
|
|
# # Auto-start when daq.py is imported/run
|
|
# if __name__ == "__main__":
|
|
# # If running daq.py directly
|
|
# start_image_stats_receiver()
|
|
#
|
|
# try:
|
|
# # Your existing daq.py code continues here
|
|
# print("DAQ system running with image statistics monitoring...")
|
|
#
|
|
# # Keep the program running
|
|
# while True:
|
|
# time.sleep(1)
|
|
#
|
|
# # Optional: Access latest stats in your main code
|
|
# latest = get_latest_image_stats()
|
|
# if latest and latest['message_count'] % 10 == 0: # Every 10th message
|
|
# print(f"Main thread sees: Mean={latest['mean']:.1f}")
|
|
#
|
|
# except KeyboardInterrupt:
|
|
# print("\nShutting down...")
|
|
# finally:
|
|
# stop_image_stats_receiver()
|
|
#
|
|
# else:
|
|
# # If daq.py is imported as a module, auto-start the receiver
|
|
# start_image_stats_receiver()
|