phase2_detect.py

"""
Phase 2 - YOLOv8 Detection + Bounding Box Visualisation
Runs real YOLOv8 detections on warehouse images, draws annotated
bounding boxes, saves results, and optionally uploads to GCS/S3.

Run:
    python phase2_detect.py --input data/sample_images/
    python phase2_detect.py --input data/sample_images/warehouse_01.jpg
    python phase2_detect.py --input data/sample_images/ --upload gcs
    python phase2_detect.py --input data/sample_images/ --upload aws
"""

import argparse
import json
from pathlib import Path
import numpy as np
import cv2
from loguru import logger
from datetime import datetime


# ─── Config ───────────────────────────────────────────────────────
OUTPUT_DIR       = Path("output/annotated")
REPORT_DIR       = Path("output/reports")
MODEL_NAME       = "yolov8n.pt"        # nano = fastest on CPU; swap to yolov8s.pt for better accuracy
CONF_THRESHOLD   = 0.35
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

# Warehouse-relevant COCO classes (YOLOv8 default model)
WAREHOUSE_LABELS = {
    "person":    ("worker",   (0, 200, 0)),       # green
    "truck":     ("vehicle",  (0, 100, 255)),      # orange
    "car":       ("vehicle",  (0, 100, 255)),
    "motorcycle":("vehicle",  (0, 100, 255)),
    "bicycle":   ("vehicle",  (0, 100, 255)),
    "backpack":  ("parcel",   (255, 200, 0)),      # cyan
    "suitcase":  ("parcel",   (255, 200, 0)),
    "bottle":    ("item",     (200, 200, 200)),
    "chair":     ("obstacle", (0, 0, 220)),        # red
    "couch":     ("obstacle", (0, 0, 220)),
    "box":       ("pallet",   (255, 100, 0)),
    "laptop":    ("equipment",(180, 0, 180)),
    "tv":        ("equipment",(180, 0, 180)),
}


def load_model():
    """Load YOLOv8 model (downloads automatically on first run ~6MB)."""
    from ultralytics import YOLO
    logger.info(f"Loading YOLOv8 model: {MODEL_NAME}")
    model = YOLO(MODEL_NAME)
    logger.success("Model loaded ✓")
    return model


def detect_image(model, image_path: Path) -> dict:
    """
    Run YOLOv8 on a single image.

    Returns:
        dict with image array, detections list, and metadata
    """
    image = cv2.imread(str(image_path))
    if image is None:
        logger.warning(f"Could not load: {image_path}")
        return None

    h, w = image.shape[:2]
    results = model(image, conf=CONF_THRESHOLD, verbose=False)

    detections = []
    for result in results:
        for box in result.boxes:
            cls_name = result.names[int(box.cls)]
            label_info = WAREHOUSE_LABELS.get(cls_name)
            warehouse_label = label_info[0] if label_info else cls_name
            colour        = label_info[1] if label_info else (200, 200, 200)
            conf          = float(box.conf)
            x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]

            detections.append({
                "original_class":  cls_name,
                "warehouse_label": warehouse_label,
                "confidence":      round(conf, 3),
                "bbox":            [x1, y1, x2, y2],
                "colour":          colour,
            })

    logger.info(f"  {image_path.name}: {len(detections)} object(s) detected")
    return {
        "image":      image,
        "path":       image_path,
        "detections": detections,
        "image_size": (w, h),
    }


def draw_annotations(result: dict) -> np.ndarray:
    """
    Draw bounding boxes and labels on the image.
    Returns annotated image as np.ndarray.
    """
    image = result["image"].copy()

    for det in result["detections"]:
        x1, y1, x2, y2 = det["bbox"]
        colour  = det["colour"]
        label   = f"{det['warehouse_label']} {det['confidence']:.0%}"

        # Bounding box
        cv2.rectangle(image, (x1, y1), (x2, y2), colour, 2)

        # Label background
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
        cv2.rectangle(image, (x1, y1 - th - 8), (x1 + tw + 4, y1), colour, -1)

        # Label text
        cv2.putText(
            image, label,
            (x1 + 2, y1 - 4),
            cv2.FONT_HERSHEY_SIMPLEX, 0.55,
            (255, 255, 255), 1, cv2.LINE_AA,
        )

    # Summary overlay (top-left)
    summary = f"Objects: {len(result['detections'])}  |  Model: {MODEL_NAME}"
    cv2.putText(image, summary, (10, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 3, cv2.LINE_AA)
    cv2.putText(image, summary, (10, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1, cv2.LINE_AA)

    return image


def save_annotated(annotated: np.ndarray, original_path: Path) -> Path:
    """Save annotated image to output directory."""
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    out_path = OUTPUT_DIR / f"annotated_{original_path.name}"
    cv2.imwrite(str(out_path), annotated)
    logger.success(f"  Saved annotated: {out_path}")
    return out_path


def save_report(results: list, output_path: Path) -> Path:
    """Save JSON detection report."""
    REPORT_DIR.mkdir(parents=True, exist_ok=True)
    report = {
        "generated_at": datetime.utcnow().isoformat(),
        "model": MODEL_NAME,
        "confidence_threshold": CONF_THRESHOLD,
        "images_processed": len(results),
        "total_detections": sum(len(r["detections"]) for r in results),
        "results": [
            {
                "image": r["path"].name,
                "image_size": r["image_size"],
                "detections": [
                    {k: v for k, v in d.items() if k != "colour"}
                    for d in r["detections"]
                ],
            }
            for r in results
        ],
    }
    with open(output_path, "w") as f:
        json.dump(report, f, indent=2)
    logger.success(f"Report saved: {output_path}")
    return output_path


def upload_to_cloud(files: list[Path], provider: str):
    """Upload annotated images and report to GCS or S3."""
    if provider == "gcs":
        from cloud_infra.setup_gcs import upload_image
        for f in files:
            upload_image(f, gcs_folder="output/phase2/")
    elif provider == "aws":
        from cloud_infra.setup_aws import upload_image
        for f in files:
            upload_image(f, s3_folder="output/phase2/")
    logger.success(f"Uploaded {len(files)} file(s) to {provider.upper()}")


def run_pipeline(input_path: Path, upload: str = None):
    """Main Phase 2 pipeline."""
    # 1. Collect image paths
    if input_path.is_file():
        image_paths = [input_path]
    elif input_path.is_dir():
        image_paths = sorted([
            p for p in input_path.iterdir()
            if p.suffix.lower() in IMAGE_EXTENSIONS
        ])
    else:
        logger.error(f"Input not found: {input_path}")
        return

    if not image_paths:
        logger.error("No images found. Run: python data/download_samples.py first")
        return

    logger.info(f"Processing {len(image_paths)} image(s) from {input_path}")

    # 2. Load model
    model = load_model()

    # 3. Detect + annotate
    all_results = []
    annotated_paths = []

    for img_path in image_paths:
        result = detect_image(model, img_path)
        if result is None:
            continue

        annotated = draw_annotations(result)
        out_path = save_annotated(annotated, img_path)

        all_results.append(result)
        annotated_paths.append(out_path)

    # 4. Save report
    report_path = REPORT_DIR / f"phase2_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    save_report(all_results, report_path)
    annotated_paths.append(report_path)

    # 5. Summary
    total = sum(len(r["detections"]) for r in all_results)
    logger.success(f"\n{'='*50}")
    logger.success(f"  Phase 2 Complete!")
    logger.success(f"  Images processed : {len(all_results)}")
    logger.success(f"  Total detections : {total}")
    logger.success(f"  Annotated images : {OUTPUT_DIR}/")
    logger.success(f"  Report           : {report_path}")
    logger.success(f"{'='*50}")

    # 6. Optional cloud upload
    if upload:
        upload_to_cloud(annotated_paths, upload)

    return all_results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Phase 2: YOLOv8 Warehouse Detection")
    parser.add_argument("--input",  type=str, default="data/sample_images/", help="Image or folder")
    parser.add_argument("--upload", type=str, choices=["gcs", "aws"], default=None, help="Upload to cloud")
    parser.add_argument("--conf",   type=float, default=CONF_THRESHOLD, help="Confidence threshold")
    args = parser.parse_args()

    CONF_THRESHOLD = args.conf
    run_pipeline(Path(args.input), upload=args.upload)