inference script added

rfdetr_seg_infer.py

@@ -0,0 +1,190 @@
+#!/usr/bin/env python3
+"""
+rfdetr_seg_infer.py
+Simple RF-DETR Segmentation inference script.
+
+Usage example:
+python rfdetr_seg_infer.py \
+    --image path/to/image.jpg \
+    --weights-url "https://huggingface.co/<user>/<repo>/resolve/main/checkpoint_best_total.pth" \
+    --out annotated.png
+
+Or provide a local weights path:
+python rfdetr_seg_infer.py --image image.jpg --weights /path/to/checkpoint_best_total.pth
+"""
+
+import argparse
+import os
+import requests
+from io import BytesIO
+
+from PIL import Image
+import numpy as np
+
+import supervision as sv
+from rfdetr import RFDETRSegPreview
+
+
+def download_file(url: str, dst: str, chunk_size: int = 8192):
+    if os.path.exists(dst):
+        return dst
+    print(f"Downloading weights from {url} -> {dst}")
+    r = requests.get(url, stream=True)
+    r.raise_for_status()
+    with open(dst, "wb") as f:
+        for chunk in r.iter_content(chunk_size=chunk_size):
+            if chunk:
+                f.write(chunk)
+    print("Download complete.")
+    return dst
+
+
+def load_image(image_path_or_url: str) -> Image.Image:
+    # Accept local path or http(s) url
+    if image_path_or_url.startswith("http://") or image_path_or_url.startswith("https://"):
+        resp = requests.get(image_path_or_url)
+        resp.raise_for_status()
+        return Image.open(BytesIO(resp.content)).convert("RGB")
+    return Image.open(image_path_or_url).convert("RGB")
+
+
+def annotate_segmentation(image: Image.Image, detections: sv.Detections, classes: dict[int, str] | None = None) -> Image.Image:
+    """
+    Annotate image with masks, polygons and labels using supervision.
+    classes: optional mapping from class_id -> class_name (int: str)
+    """
+    # Color palette (expandable)
+    palette = sv.ColorPalette.from_hex([
+        "#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff",
+        "#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00",
+    ])
+
+    text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
+    mask_annotator = sv.MaskAnnotator(color=palette)
+    polygon_annotator = sv.PolygonAnnotator(color=sv.Color.WHITE)
+    label_annotator = sv.LabelAnnotator(
+        color=palette,
+        text_color=sv.Color.BLACK,
+        text_scale=text_scale,
+        text_position=sv.Position.CENTER_OF_MASS
+    )
+
+    labels = [
+        f"{(classes.get(class_id) if classes else str(class_id))} {conf:.2f}"
+        for class_id, conf in zip(detections.class_id, detections.confidence)
+    ]
+
+    out = image.copy()
+    out = mask_annotator.annotate(out, detections)
+    out = polygon_annotator.annotate(out, detections)
+    out = label_annotator.annotate(out, detections, labels)
+    return out
+
+
+def save_combined_color_mask(detections: sv.Detections, out_path: str, image_size):
+    """
+    Create and save a combined per-instance color mask (RGB) where each instance gets a unique color.
+    If detections has `masks` attribute (list/array of binary masks), uses that.
+    """
+    masks = getattr(detections, "masks", None)
+    if masks is None:
+        print("[INFO] No masks available on detections; skipping combined mask save.")
+        return False
+
+    # masks expected shape: (N, H, W) or list of (H, W) masks
+    if isinstance(masks, list):
+        masks_arr = np.stack([np.asarray(m, dtype=bool) for m in masks], axis=0)
+    else:
+        masks_arr = np.asarray(masks)
+        # if (H,W,N) transpose is sometimes used; try to normalize
+        if masks_arr.ndim == 3 and masks_arr.shape[0] == image_size[1] and masks_arr.shape[1] == image_size[0]:
+            # shape likely (H, W, N) -> transpose to (N, H, W)
+            masks_arr = masks_arr.transpose(2, 0, 1)
+
+    H, W = image_size[1], image_size[0]
+    combined = np.zeros((H, W, 3), dtype=np.uint8)
+
+    palette = sv.ColorPalette.from_hex([
+        "#FF0000", "#00FF00", "#0000FF", "#FFFF00", "#FF00FF", "#00FFFF",
+        "#800000", "#008000", "#000080", "#808000", "#800080", "#008080"
+    ])
+    palette_rgb = [tuple(int(h.lstrip("#")[i:i+2], 16) for i in (0, 2, 4)) for h in palette.hex_list]
+
+    for i in range(masks_arr.shape[0]):
+        mask = masks_arr[i].astype(bool)
+        color = palette_rgb[i % len(palette_rgb)]
+        # paint color where mask is True, simple overwrite (later instances will overwrite earlier)
+        combined[mask] = color
+
+    Image.fromarray(combined).save(out_path)
+    print(f"[INFO] Saved combined color mask to: {out_path}")
+    return True
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", required=True, help="Path or URL to input image")
+    parser.add_argument("--weights-url", help="HF resolve URL to checkpoint .pth (will be downloaded)")
+    parser.add_argument("--weights", help="Local path to checkpoint .pth")
+    parser.add_argument("--threshold", type=float, default=0.3, help="Confidence threshold")
+    parser.add_argument("--out", default="annotated.png", help="Path to save annotated overlay image")
+    parser.add_argument("--save-mask", default=None, help="Path to save combined color mask PNG (optional)")
+    parser.add_argument("--no-optimize", action="store_true", help="Skip optimize_for_inference()")
+    args = parser.parse_args()
+
+    # resolve weights: prioritize local --weights, then --weights-url
+    weights_path = None
+    if args.weights:
+        if not os.path.exists(args.weights):
+            raise FileNotFoundError(f"Weights file not found: {args.weights}")
+        weights_path = args.weights
+    elif args.weights_url:
+        # pick filename from url last segment
+        fname = os.path.basename(args.weights_url.split("?")[0])
+        if not fname:
+            fname = "checkpoint_best_total.pth"
+        weights_path = os.path.abspath(fname)
+        if not os.path.exists(weights_path):
+            download_file(args.weights_url, weights_path)
+    else:
+        weights_path = None  # will use internal default if RFC allows
+
+    # load image (PIL)
+    image = load_image(args.image)
+    print(f"[INFO] Image loaded: {args.image} (size={image.size})")
+
+    # instantiate model
+    if weights_path:
+        model = RFDETRSegPreview(pretrain_weights=weights_path)
+    else:
+        model = RFDETRSegPreview()
+
+    if not args.no_optimize:
+        try:
+            model.optimize_for_inference()
+        except Exception as e:
+            print(f"[WARN] optimize_for_inference() failed or skipped: {e}")
+
+    # model.predict accepts PIL.Image or path
+    detections = model.predict(image, threshold=args.threshold)
+    print("[INFO] Inference done. Found", len(detections.boxes) if hasattr(detections, "boxes") else getattr(detections, "xyxy", "N/A"))
+
+    # annotate
+    # If you trained on custom classes the model should include class mapping. If not,
+    # you can pass classes dict e.g. {0: "Tulsi", 1: "Leaf"} depending on your dataset.
+    annotated = annotate_segmentation(image, detections, classes=None)
+    annotated.save(args.out)
+    print(f"[INFO] Saved annotated image to: {args.out}")
+
+    # optional combined color mask (per-instance overlay)
+    if args.save_mask:
+        ok = save_combined_color_mask(detections, args.save_mask, image_size=image.size)
+        if not ok:
+            print("[INFO] Combined mask not created (no mask data).")
+
+    # Also return annotated PIL object for interactive sessions
+    return annotated
+
+
+if __name__ == "__main__":
+    main()