Spaces:
Running
Running
File size: 5,086 Bytes
243f893 c1f2912 243f893 c1f2912 243f893 5b4cc5f 243f893 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import os
import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download
# Retrieve the token from environment variables
hf_token = os.environ.get("HF_TOKEN")
# Download the private model.onnx file
model_path = hf_hub_download(
repo_id="ymp5078/RulerNet",
filename="model.onnx",
use_auth_token=hf_token
)
# ---- Load ONNX Model ----
ort_session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
# ---- Utility Function ----
def outward_cumsum(initial_point, line_direction, directions, n):
left_directions = directions[n < 0][::-1]
right_directions = directions[n >= 0]
left_increments = -np.expand_dims(left_directions, axis=1) * line_direction
right_increments = np.expand_dims(right_directions, axis=1) * line_direction
zero = np.zeros((1, 2), dtype=initial_point.dtype)
left_cumulative = np.cumsum(np.vstack([zero, left_increments]), axis=0)
right_cumulative = np.cumsum(np.vstack([zero, right_increments]), axis=0)
left_points = initial_point + left_cumulative
right_points = initial_point + right_cumulative
extended_points = np.vstack([left_points[::-1], right_points[1:]])
return extended_points
# ---- Image Preprocessing ----
def preprocess_image(pil_img):
image = np.array(pil_img).astype(np.float32) / 255.0
if image.ndim == 2:
image = np.stack([image]*3, axis=-1)
elif image.shape[2] == 4:
image = image[:, :, :3]
resized = np.zeros((768, 768, 3), dtype=np.float32)
h, w = image.shape[:2]
scale = min(768 / w, 768 / h)
new_w, new_h = int(w * scale), int(h * scale)
image_resized = np.array(Image.fromarray((image * 255).astype(np.uint8)).resize((new_w, new_h))).astype(np.float32) / 255.0
top = (768 - new_h) // 2
left = (768 - new_w) // 2
resized[top:top+new_h, left:left+new_w] = image_resized
input_tensor = np.transpose(resized, (2, 0, 1))[np.newaxis, ...].copy()
return input_tensor, (top, left, new_h, new_w)
# ---- Main Inference and Drawing ----
def infer_and_draw(image_pil):
image_tensor, _ = preprocess_image(image_pil)
ort_inputs = {"input": image_tensor}
ort_outs = ort_session.run(None, ort_inputs)
left_point_2d_reconstructed = ort_outs[0][0] # shape: (2,)
dist = ort_outs[1][0][0] # scalar
ratio = ort_outs[2][0][0] # scalar
direction = ort_outs[3][0] # shape: (2,)
points_info = ort_outs[4][0]
min_x, min_y, max_x, max_y = points_info[1:].tolist()
num_points = int(points_info[0])
n = np.arange(-num_points, num_points + 1)
directions = (ratio ** n) * dist
extended_points = outward_cumsum(left_point_2d_reconstructed, direction, directions, n)
within_bounds = (
(extended_points[:, 0] >= min_x) & (extended_points[:, 0] <= max_x) &
(extended_points[:, 1] >= min_y) & (extended_points[:, 1] <= max_y)
)
best_generated_points = extended_points[within_bounds]
if len(best_generated_points) > 1:
diffs = np.linalg.norm(best_generated_points[:-1] - best_generated_points[1:], axis=1)
pred_pix_cm = np.nanmedian(diffs)
else:
pred_pix_cm = 0.0
# ---- Draw on image ----
np_img = (np.transpose(image_tensor[0], (1, 2, 0)) * 255).astype(np.uint8)
pil_img = Image.fromarray(np_img)
draw = ImageDraw.Draw(pil_img)
# Pixel grid overlay
grid_spacing = 50
width, height = pil_img.size
for x in range(0, width, grid_spacing):
draw.line([(x, 0), (x, height)], fill=(200, 200, 200), width=1)
for y in range(0, height, grid_spacing):
draw.line([(0, y), (width, y)], fill=(200, 200, 200), width=1)
r = 3
for x, y in best_generated_points:
draw.ellipse((x - r, y - r, x + r, y + r), fill="red")
text = f"Pix/cm: {pred_pix_cm:.2f}"
text_position = (10, 10)
text_size = draw.textbbox(text_position, text)
padding = 4
rect_coords = (
text_size[0] - padding,
text_size[1] - padding,
text_size[2] + padding,
text_size[3] + padding
)
draw.rectangle(rect_coords, fill="white")
draw.text(text_position, text, fill="black")
return pil_img, f"{pred_pix_cm:.3f}", f"{ratio:.3f}"
if __name__ == '__main__':
demo = gr.Interface(
fn=infer_and_draw,
inputs=gr.Image(type="pil"),
outputs=[
gr.Image(label="Generated Ruler Points"),
gr.Textbox(label="Predicted median pixel-per-centimeter value (with grid lines every 50 pixels):"),
gr.Textbox(label="Predicted geometric progression ratio:")
],
examples=[
["examples/sample1.jpg"],
["examples/sample2.jpg"],
["examples/sample3.jpg"],
["examples/sample4.jpg"],
["examples/sample5.jpg"],
["examples/sample6.jpg"]
],
title="ONNX (CPU) Ruler Model Visualizer",
description="Upload an image to visualize ruler points generated by the ONNX model."
)
demo.launch()
|