yamero999 commited on
Commit
091cd2d
·
verified ·
1 Parent(s): adf16b0

Upload example_usage.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. 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()