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