hatch-finder-3.5m

GreenMap/hatch-finder-3.5m

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

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