add app file

app.py

@@ -0,0 +1,137 @@
+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")
+
+    x0, y0 = left_point_2d_reconstructed[0]
+    draw.ellipse((x0 - r, y0 - r, x0 + r, y0 + r), fill="blue")
+
+    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:.2f}"
+
+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):")
+        ],
+        title="ONNX (CPU) Ruler Model Visualizer",
+        description="Upload an image to visualize ruler points generated by the ONNX model."
+    )
+
+    demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+numpy
+onnxruntime
+Pillow
+huggingface_hub