File size: 2,782 Bytes
091cd2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
"""
Example usage of Ultimate V2 Chess Board Segmentation ONNX model
"""
import onnxruntime as ort
import numpy as np
import cv2
import matplotlib.pyplot as plt
def load_model(model_path):
"""Load ONNX model"""
session = ort.InferenceSession(model_path)
return session
def preprocess_image(image_path):
"""Preprocess image for model input"""
# Load image
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Resize to model input size
image_resized = cv2.resize(image_rgb, (256, 256))
# Normalize to [0, 1]
image_normalized = image_resized.astype(np.float32) / 255.0
# Convert to model input format (NCHW)
input_tensor = np.transpose(image_normalized, (2, 0, 1))[np.newaxis, ...]
return input_tensor, image_rgb
def run_inference(session, input_tensor):
"""Run inference on the model"""
# Get input name
input_name = session.get_inputs()[0].name
# Run inference
outputs = session.run(None, {input_name: input_tensor})
# Apply sigmoid to get probabilities
mask = 1.0 / (1.0 + np.exp(-outputs[0]))
return mask.squeeze()
def visualize_results(original_image, mask, threshold=0.5):
"""Visualize the segmentation results"""
# Create binary mask
binary_mask = (mask > threshold).astype(np.uint8) * 255
# Create overlay
overlay = original_image.copy()
overlay[binary_mask > 0] = [255, 0, 0] # Red overlay
# Plot results
fig, axes = plt.subplots(1, 4, figsize=(16, 4))
axes[0].imshow(original_image)
axes[0].set_title('Original Image')
axes[0].axis('off')
axes[1].imshow(mask, cmap='hot')
axes[1].set_title('Segmentation Heatmap')
axes[1].axis('off')
axes[2].imshow(binary_mask, cmap='gray')
axes[2].set_title('Binary Mask')
axes[2].axis('off')
axes[3].imshow(overlay)
axes[3].set_title('Overlay')
axes[3].axis('off')
plt.tight_layout()
plt.show()
def main():
"""Main example function"""
# Load model
model_path = "ultimate_v2_breakthrough_accurate.onnx"
session = load_model(model_path)
# Process image
image_path = "chess_board.jpg" # Replace with your image
input_tensor, original_image = preprocess_image(image_path)
# Run inference
mask = run_inference(session, input_tensor)
# Visualize results
visualize_results(original_image, mask)
print(f"✅ Chess board segmentation completed!")
print(f"📊 Mask shape: {mask.shape}")
print(f"📈 Mask range: {mask.min():.3f} - {mask.max():.3f}")
if __name__ == "__main__":
main()
|