signature/inference.py

import cv2
import os
import sys
import json
import supervision as sv
from huggingface_hub import hf_hub_download, login
from ultralytics import YOLO
from pathlib import Path


def detect_signatures(image_path, model=None, output_dir=None, signatures_dir=None, save_crops=True):
    """
    Detect signatures in a single image.

    Args:
        image_path: Path to the input image
        model: YOLO model instance (if None, will load/create one)
        output_dir: Directory for output files (optional)
        signatures_dir: Directory for cropped signatures (optional)
        save_crops: Whether to save cropped signature images

    Returns:
        dict: Detection results with structure:
            {
                "image": image_filename,
                "image_width": int,
                "image_height": int,
                "signatures": [...]
            }
    """
    # Load model if not provided
    if model is None:
        local_model_path = Path("yolov8s.pt")
        if local_model_path.exists():
            model_path = str(local_model_path)
        else:
            try:
                # Get HF token from environment (for gated models)
                hf_token = os.environ.get(
                    "HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
                model_path = hf_hub_download(
                    repo_id="tech4humans/yolov8s-signature-detector",
                    filename="yolov8s.pt",
                    token=hf_token  # Pass token for gated repos
                )
            except Exception as e:
                raise RuntimeError(f"Failed to load signature model: {e}")
        model = YOLO(model_path)

    # Set up paths (only if we need to save crops)
    image_file = Path(image_path)
    if save_crops:
        if output_dir is None:
            output_dir = Path("outputs")
        else:
            output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True)

        if signatures_dir is None:
            signatures_dir = output_dir / "signatures"
        else:
            signatures_dir = Path(signatures_dir)
        signatures_dir.mkdir(exist_ok=True)
    else:
        # Dummy paths when not saving
        output_dir = None
        signatures_dir = None

    # Read image
    image = cv2.imread(str(image_path))
    if image is None:
        raise ValueError(f"Could not read image: {image_path}")

    # Get image dimensions
    image_height, image_width = image.shape[:2]

    # Run inference
    results = model(str(image_path))
    detections = sv.Detections.from_ultralytics(results[0])

    # Store detection data
    image_detections = {
        "image": image_file.name,
        "image_width": int(image_width),
        "image_height": int(image_height),
        "signatures": []
    }

    # Process detections
    if len(detections) > 0:
        for i, (xyxy, confidence, class_id) in enumerate(zip(
            detections.xyxy, detections.confidence, detections.class_id
        )):
            x1, y1, x2, y2 = xyxy

            # Store detection data
            detection_data = {
                "signature_id": i + 1,
                "confidence": float(confidence),
                "bbox": {
                    "x1": float(x1),
                    "y1": float(y1),
                    "x2": float(x2),
                    "y2": float(y2),
                    "width": float(x2 - x1),
                    "height": float(y2 - y1)
                },
                "class_id": int(class_id)
            }

            # Crop and save individual signature if requested
            if save_crops and signatures_dir is not None:
                x1_int, y1_int, x2_int, y2_int = int(
                    x1), int(y1), int(x2), int(y2)
                x1_int = max(0, x1_int)
                y1_int = max(0, y1_int)
                x2_int = min(image.shape[1], x2_int)
                y2_int = min(image.shape[0], y2_int)

                signature_crop = image[y1_int:y2_int, x1_int:x2_int]
                signature_filename = f"{image_file.stem}_signature_{i+1}.jpg"
                signature_path = signatures_dir / signature_filename
                cv2.imwrite(str(signature_path), signature_crop)
                detection_data["cropped_path"] = str(signature_path)

            image_detections["signatures"].append(detection_data)

    return image_detections


def main():
    # Check if model file exists locally first
    local_model_path = Path("yolov8s.pt")

    if local_model_path.exists():
        print(f"Using local model file: {local_model_path}", flush=True)
        model_path = str(local_model_path)
    else:
        # Try to download model from Hugging Face
        print("Downloading model from Hugging Face...", flush=True)
        try:
            model_path = hf_hub_download(
                repo_id="tech4humans/yolov8s-signature-detector",
                filename="yolov8s.pt"
            )
        except Exception as e:
            if "401" in str(e) or "GatedRepoError" in str(type(e).__name__) or "Unauthorized" in str(e):
                print("\n" + "="*70)
                print("ERROR: Authentication required to access this model.")
                print("="*70)
                print(
                    "\nThis repository is gated and requires Hugging Face authentication.")
                print("\nTo authenticate, run one of the following:")
                print("  1. huggingface-cli login")
                print("  2. Or set your token: export HF_TOKEN=your_token_here")
                print("\nAfter authentication, run this script again.")
                print("="*70)
                sys.exit(1)
            else:
                print(f"\nError downloading model: {e}")
                print("\nYou can also download the model manually:")
                print(
                    "  huggingface-cli download tech4humans/yolov8s-signature-detector yolov8s.pt")
                print("\nOr place yolov8s.pt in the current directory.")
                sys.exit(1)

    # Load the model
    print("Loading model...")
    model = YOLO(model_path)

    # Set up paths
    input_dir = Path("inputs")
    output_dir = Path("outputs")
    signatures_dir = output_dir / "signatures"  # Directory for cropped signatures
    output_dir.mkdir(exist_ok=True)
    signatures_dir.mkdir(exist_ok=True)

    # Store all detections for JSON export
    all_detections = []

    # Get all image files from inputs directory
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    image_files = [f for f in input_dir.iterdir()
                   if f.suffix.lower() in image_extensions]

    if not image_files:
        print(f"No images found in {input_dir}/")
        return

    print(f"Found {len(image_files)} image(s) to process")

    # Process each image
    box_annotator = sv.BoxAnnotator()

    for image_file in image_files:
        print(f"\nProcessing: {image_file.name}")

        try:
            # Use the reusable function
            image_detections = detect_signatures(
                str(image_file),
                model=model,
                output_dir=output_dir,
                signatures_dir=signatures_dir,
                save_crops=True
            )

            # Read image for annotation
            image = cv2.imread(str(image_file))
            results = model(str(image_file))
            detections = sv.Detections.from_ultralytics(results[0])

            if len(detections) > 0:
                print(f"  Found {len(detections)} signature(s)")
                for i, sig in enumerate(image_detections["signatures"]):
                    bbox = sig["bbox"]
                    print(
                        f"    Signature {i+1}: confidence={sig['confidence']:.2f}, bbox=[{bbox['x1']:.1f}, {bbox['y1']:.1f}, {bbox['x2']:.1f}, {bbox['y2']:.1f}]")
                    if "cropped_path" in sig:
                        print(
                            f"      Saved cropped signature to: {sig['cropped_path']}")
            else:
                print("  No signatures detected")

            all_detections.append(image_detections)

            # Annotate image with bounding boxes
            annotated_image = box_annotator.annotate(
                scene=image.copy(),
                detections=detections
            )

            # Save annotated image
            output_path = output_dir / f"detected_{image_file.name}"
            cv2.imwrite(str(output_path), annotated_image)
            print(f"  Saved annotated image to: {output_path}")
        except Exception as e:
            print(f"  Error processing {image_file.name}: {str(e)}")
            continue

    # Save all coordinates to JSON file
    json_path = output_dir / "signature_coordinates.json"
    with open(json_path, 'w') as f:
        json.dump(all_detections, f, indent=2)
    print(f"\n{'='*70}")
    print(f"Saved all signature coordinates to: {json_path}")
    print(f"{'='*70}")


if __name__ == "__main__":
    main()