GreenMap/hatch-finder-3.5m

hf_preview

HatchFinder is a lightweight reference-conditioned model for locating a given hatch pattern in architectural and construction drawings.

The model takes a blueprint, a search mask defining the region of interest, and a reference image of the hatch pattern to find. It produces a dense pixel-wise heatmap indicating where the reference hatch is present.

GitHub: GreenMap-chan/hatch-finder

Overview

HatchFinder is designed for reference-based hatch detection rather than classification into a fixed set of hatch classes.

  • Finds a hatch pattern provided as a reference image.
  • Produces a dense pixel-wise probability heatmap.
  • Uses a search mask to restrict detection to selected regions of the drawing.
  • Does not require the target hatch to belong to a predefined class.
  • Designed for architectural plans, construction drawings, and similar technical documents.
  • Approximately 3.5M parameters.

Inputs

The model takes three inputs:

  • Drawing โ€” RGB image of the architectural or construction drawing.
  • Search mask โ€” binary mask specifying the region in which the hatch should be searched for.
  • Reference hatch โ€” RGB image containing an example of the hatch pattern to locate.

Since the model was trained on 896ร—896 images, I strongly recommend using the same input size for Drawing.

Output

The model outputs a single-channel dense logits map with the same spatial resolution as the drawing.

After applying sigmoid, each pixel represents the predicted probability that it belongs to the reference hatch pattern.

heatmap = torch.sigmoid(logits)
heatmap = heatmap * search_mask

Architecture

HatchFinder uses a custom multi-scale convolutional architecture consisting of:

  • Multi-scale drawing encoder.
  • Multi-scale reference hatch encoder.
  • Learned drawing-to-reference matching features.
  • Gated multi-scale matching.
  • FiLM-based reference conditioning.
  • U-Net-like heatmap decoder.

The model is trained end-to-end for reference-conditioned pixel-level hatch detection.

Validation Metrics

Metric Value
Validation loss 0.05881
BCE loss 0.01343
Dice loss 0.05672
Dice coefficient 0.94328

The metrics above were measured on the validation split used for this model and should not be interpreted as performance on arbitrary blueprint datasets.

Training

  • Task: Reference-conditioned hatch segmentation
  • Framework: PyTorch
  • Parameters: ~3.5M
  • Loss: Masked BCEWithLogits + Dice loss
  • Training examples: 7,000
  • Input: RGB drawing + binary search mask + RGB reference hatch
  • Input drawing size: 896 x 896
  • Output: Pixel-wise hatch logits

Training includes geometric and visual augmentations applied independently where appropriate to the drawing and reference hatch.

Installation

pip install hatchfinder huggingface_hub

Inference

The following example is self-contained: it downloads the model and one input triplet from this repository, runs inference, and saves both the probability heatmap and a thresholded prediction. It also creates a visualization in output/0000435_debug.png, where predictions above confidence are overlaid in red and the area outside the search mask is darkened.

from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms.functional import to_pil_image

from hatchfinder import HatchFinder

REPO_ID = "GreenMap/hatch-finder-3.5m"
SAMPLE_ID = "0000435"
OUTPUT_DIR = Path("output")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


def download(filename: str) -> Path:
    return Path(hf_hub_download(repo_id=REPO_ID, filename=filename))


# Download the weights and a drawing/search-mask/reference-hatch triplet.
model_path = download("model.pt")
drawing_path = download(f"sample_dataset/valid/drawing/{SAMPLE_ID}.png")
mask_path = download(f"sample_dataset/valid/search_mask/{SAMPLE_ID}.png")
hatch_path = download(f"sample_dataset/valid/hatch/{SAMPLE_ID}.png")

# "auto" selects CUDA when it is available and otherwise uses the CPU.
model = HatchFinder(load_model_path=model_path, device="auto")

heatmap = model.infer(
    drawing=drawing_path,
    mask=mask_path,
    hatch=hatch_path,
    debug_path=OUTPUT_DIR,
    confidence=0.5,
)

# infer() returns probabilities with shape [1, 1, height, width]. Pixels
# outside the search mask are already set to zero.
print(heatmap.shape, heatmap.min().item(), heatmap.max().item())

heatmap_cpu = heatmap[0].detach().cpu().clamp(0, 1)
to_pil_image(heatmap_cpu).save(OUTPUT_DIR / f"{SAMPLE_ID}_heatmap.png")

prediction = (heatmap_cpu >= 0.5).to(torch.uint8) * 255
Image.fromarray(prediction[0].numpy()).save(
    OUTPUT_DIR / f"{SAMPLE_ID}_prediction.png"
)

drawing and hatch can be paths or Pillow images; they are converted to RGB. mask can also be a path or a Pillow image; it is converted to grayscale and binarized at 0.5. The drawing and search mask must have the same dimensions. For best results, use a drawing size of 896 x 896, matching the training data.

See the hatch-finder repository for the source code, training configuration, and further examples.

License

Apache License 2.0

Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Collection including GreenMap/hatch-finder-3.5m