test.py

import os
import io
import base64
import json
import torch
import numpy as np
from dataclasses import dataclass, asdict
from typing import List, Dict, Any, Optional
from pathlib import Path

from flask import Flask, request, render_template_string
from PIL import Image

# --- Machine Learning Imports ---
try:
    import supervision as sv
    from groundingdino.util.inference import Model as DinoModel
    from sam2.build_sam import build_sam2
    from sam2.sam2_image_predictor import SAM2ImagePredictor
except ImportError as e:
    print("WARNING: ML libraries not found. Ensure 'groundingdino', 'sam2', 'supervision', and 'torch' are installed.")
    raise e

# =============================================================
# Configuration & Constants
# =============================================================

ALLOWED_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}

# DEVICE SETUP
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Running on device: {DEVICE}")

# =============================================================
# PATHS (YOUR CONFIGURATION)
# =============================================================

DINO_CONFIG_PATH = r"C:\yuxuan\GroundingDINO\groundingdino\config\GroundingDINO_SwinT_OGC.py"
DINO_CHECKPOINT_PATH = r"C:\yuxuan\GroundingDINO\groundingdino_swint_ogc.pth"

SAM2_CONFIG_PATH = r"C:\yuxuan\GroundingDINO\sam2_repo\sam2\configs\sam2\sam2_hiera_l.yaml"
SAM2_CHECKPOINT_PATH = r"C:\yuxuan\GroundingDINO\checkpoints\sam2_hiera_large.pt"

# =============================================================
# Global Model Loading
# =============================================================

print("⏳ Loading Grounding DINO...")
if not os.path.exists(DINO_CHECKPOINT_PATH):
    print(f"❌ Critical Error: DINO Checkpoint not found at: {os.path.abspath(DINO_CHECKPOINT_PATH)}")
    dino_model = None
else:
    try:
        dino_model = DinoModel(
            model_config_path=DINO_CONFIG_PATH,
            model_checkpoint_path=DINO_CHECKPOINT_PATH,
            device=DEVICE
        )
        print("✅ Grounding DINO loaded.")
    except Exception as e:
        print(f"❌ Error loading Grounding DINO: {e}")
        dino_model = None


print("⏳ Loading SAM 2...")
if not os.path.exists(SAM2_CHECKPOINT_PATH):
    print(f"❌ Critical Error: SAM 2 Checkpoint not found at: {os.path.abspath(SAM2_CHECKPOINT_PATH)}")
    sam2_predictor = None
else:
    try:
        sam2_model = build_sam2(SAM2_CONFIG_PATH, SAM2_CHECKPOINT_PATH, device=DEVICE)
        sam2_predictor = SAM2ImagePredictor(sam2_model)
        print("✅ SAM 2 loaded.")
    except Exception as e:
        print(f"❌ Error loading SAM 2: {e}")
        sam2_predictor = None


# =============================================================
# Data Structures
# =============================================================

@dataclass
class Annotation:
    label: str
    score: float
    bbox: List[float]
    mask: Optional[np.ndarray] = None
    
    def __post_init__(self):
        pass

@dataclass
class ProcessedImage:
    filename: str
    image_b64: str
    annotations: List[Annotation]
    is_annotated: bool
    image_size: tuple  # (width, height)
    
    def to_export_dict(self):
        """Convert to dict for export, excluding non-serializable fields"""
        return {
            'filename': self.filename,
            'is_annotated': self.is_annotated,
            'image_size': self.image_size,
            'annotations': [
                {
                    'label': ann.label,
                    'score': ann.score,
                    'bbox': ann.bbox,
                    'mask': ann.mask  # Keep mask for export processing
                }
                for ann in self.annotations
            ]
        }

# =============================================================
# Inference Logic
# =============================================================

def run_grounded_dino_sam(image: Image.Image, prompt: str) -> Dict[str, Any]:
    """
    Runs Grounding DINO to detect boxes -> SAM 2 to generate masks.
    """
    if not dino_model or not sam2_predictor:
        return {"annotated_image": image, "annotations": [], "detections": None}

    with torch.inference_mode():
        image = image.convert("RGB")
        image_np = np.array(image)
        image_size = image.size  # (width, height)

        # 1. Grounding DINO
        detections = dino_model.predict_with_classes(
            image=image_np,
            classes=[prompt],
            box_threshold=0.35,
            text_threshold=0.25
        )

        if len(detections.xyxy) == 0:
            return {"annotated_image": image, "annotations": [], "detections": None}

        # Fix class_id
        if detections.class_id is None:
            detections.class_id = np.zeros(len(detections.xyxy), dtype=int)
        else:
            try:
                detections.class_id = detections.class_id.astype(int)
            except:
                detections.class_id = np.zeros(len(detections.xyxy), dtype=int)

        # 2. SAM 2
        sam2_predictor.set_image(image_np)
        masks, scores, _ = sam2_predictor.predict(
            point_coords=None,
            point_labels=None,
            box=detections.xyxy,
            multimask_output=False
        )
        if masks.ndim == 4:
            masks = masks.squeeze(1)
        detections.mask = masks > 0.0

        # 3. Visualization
        box_annotator = sv.BoxAnnotator()
        annotated_frame = box_annotator.annotate(scene=image_np.copy(), detections=detections)
        
        mask_annotator = sv.MaskAnnotator()
        annotated_frame = mask_annotator.annotate(scene=annotated_frame, detections=detections)

        # 4. Serialize
        anns = []
        for xyxy, mask, confidence, class_id, tracker_id, *_ in detections:
            anns.append(Annotation(
                label=prompt,
                score=float(confidence) if confidence else 0.0,
                bbox=[float(x) for x in xyxy],
                mask=mask.astype(np.uint8) if mask is not None else None
            ))

        return {
            "annotated_image": Image.fromarray(annotated_frame),
            "annotations": anns,
            "detections": detections,
            "image_size": image_size
        }

def pil_to_base64(img: Image.Image) -> str:
    img.thumbnail((800, 800))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    buf.seek(0)
    return base64.b64encode(buf.read()).decode("utf-8")

def mask_to_coco_rle(mask: np.ndarray) -> Dict:
    """Convert binary mask to COCO RLE format"""
    from pycocotools import mask as coco_mask
    
    mask_uint8 = (mask * 255).astype(np.uint8)
    fortran_mask = np.asfortranarray(mask_uint8)
    rle = coco_mask.encode(fortran_mask)
    rle['counts'] = rle['counts'].decode('utf-8')
    return rle

def mask_to_polygon(mask: np.ndarray) -> List[List[float]]:
    """Convert binary mask to polygon coordinates"""
    from cv2 import findContours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE
    import cv2
    
    mask_uint8 = (mask * 255).astype(np.uint8)
    contours, _ = findContours(mask_uint8, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)
    
    polygons = []
    for contour in contours:
        if len(contour) >= 3:
            polygon = contour.flatten().tolist()
            if len(polygon) >= 6:
                polygons.append(polygon)
    
    return polygons if polygons else [[]]

def bbox_to_yolo(bbox: List[float], img_width: float, img_height: float) -> List[float]:
    """Convert [x1, y1, x2, y2] to YOLO format [center_x, center_y, width, height] normalized"""
    x1, y1, x2, y2 = bbox
    center_x = (x1 + x2) / 2 / img_width
    center_y = (y1 + y2) / 2 / img_height
    width = (x2 - x1) / img_width
    height = (y2 - y1) / img_height
    return [center_x, center_y, width, height]

def mask_to_yolo_segmentation(mask: np.ndarray, img_width: float, img_height: float) -> List[float]:
    """Convert mask to YOLO segmentation format (normalized polygon)"""
    from cv2 import findContours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE
    import cv2
    
    mask_uint8 = (mask * 255).astype(np.uint8)
    contours, _ = findContours(mask_uint8, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)
    
    if contours:
        contour = max(contours, key=cv2.contourArea)
        polygon = contour.flatten().tolist()
        # Normalize coordinates
        normalized = [p / (img_width if i % 2 == 0 else img_height) for i, p in enumerate(polygon)]
        return normalized
    return []

# =============================================================
# Export Functions
# =============================================================

def export_labels(folder: str, images_data: List[Dict], export_format: str, export_type: str, prompt: str):
    """
    Export labels in COCO or YOLO format.
    export_format: 'coco' or 'yolo'
    export_type: 'detection' or 'segmentation'
    """
    parent_dir = os.path.dirname(folder)
    labels_folder = os.path.join(parent_dir, "labels")
    
    print(f"   Folder: {folder}")
    print(f"   Parent dir: {parent_dir}")
    print(f"   Labels folder target: {labels_folder}")
    
    try:
        os.makedirs(labels_folder, exist_ok=True)
        print(f"   ✅ Labels folder created/exists: {labels_folder}")
    except Exception as e:
        print(f"   ❌ Failed to create labels folder: {e}")
        raise
    
    print(f"   Images to export: {len(images_data)}")
    
    if export_format == "coco":
        export_coco(labels_folder, images_data, export_type, prompt)
    elif export_format == "yolo":
        export_yolo(labels_folder, images_data, export_type, prompt)

def export_coco(labels_folder: str, images_data: List[Dict], export_type: str, prompt: str):
    """Export in COCO format to a single JSON file"""
    print(f"   🔹 COCO Export - Type: {export_type}")
    print(f"      Total images: {len(images_data)}")
    
    coco_data = {
        "info": {
            "description": f"Grounded SAM 2 {export_type} annotations",
            "version": "1.0"
        },
        "licenses": [],
        "images": [],
        "annotations": [],
        "categories": [{"id": 0, "name": prompt}]
    }
    
    annotation_id = 1
    total_annotations = 0
    
    for img_idx, img_data in enumerate(images_data):
        if not img_data['is_annotated'] or not img_data['annotations']:
            print(f"      [{img_idx}] {img_data['filename']}: SKIPPED (no annotations)")
            continue
        
        print(f"      [{img_idx}] {img_data['filename']}: {len(img_data['annotations'])} objects")
        
        img_entry = {
            "id": img_idx,
            "file_name": img_data['filename'],
            "width": img_data['image_size'][0],
            "height": img_data['image_size'][1]
        }
        coco_data["images"].append(img_entry)
        
        for ann_idx, ann in enumerate(img_data['annotations']):
            bbox = ann['bbox']
            x1, y1, x2, y2 = bbox
            width = x2 - x1
            height = y2 - y1
            
            print(f"        Annotation {ann_idx}: bbox={bbox}, has_mask={ann['mask'] is not None}")
            
            ann_entry = {
                "id": annotation_id,
                "image_id": img_idx,
                "category_id": 0,
                "bbox": [x1, y1, width, height],
                "area": width * height,
                "iscrowd": 0
            }
            
            if export_type == "segmentation" and ann['mask'] is not None:
                polygons = mask_to_polygon(ann['mask'])
                ann_entry["segmentation"] = polygons
                print(f"          Added segmentation with {len(polygons)} polygons")
            
            coco_data["annotations"].append(ann_entry)
            total_annotations += 1
            annotation_id += 1
    
    output_file = os.path.join(labels_folder, "annotations.json")
    with open(output_file, 'w') as f:
        json.dump(coco_data, f, indent=2)
    
    print(f"      ✅ Saved: {output_file}")
    print(f"      Total: {len(coco_data['images'])} images, {total_annotations} annotations")
    return output_file

def export_yolo(labels_folder: str, images_data: List[Dict], export_type: str, prompt: str):
    """Export in YOLO format with one .txt file per image"""
    print(f"   🔹 YOLO Export - Type: {export_type}")
    
    os.makedirs(labels_folder, exist_ok=True)
    
    exported_count = 0
    for img_data in images_data:
        if not img_data['is_annotated'] or not img_data['annotations']:
            continue
        
        # Create label file with same name as image but .txt extension
        base_name = os.path.splitext(img_data['filename'])[0]
        label_file = os.path.join(labels_folder, f"{base_name}.txt")
        
        print(f"      Processing: {img_data['filename']} -> {base_name}.txt ({len(img_data['annotations'])} annotations)")
        
        img_width, img_height = img_data['image_size']
        
        try:
            with open(label_file, 'w') as f:
                for ann in img_data['annotations']:
                    if export_type == "detection":
                        # Format: class_id center_x center_y width height
                        bbox = ann['bbox']
                        yolo_bbox = bbox_to_yolo(bbox, img_width, img_height)
                        line = f"0 {' '.join(map(str, yolo_bbox))}\n"
                        f.write(line)
                        exported_count += 1
                    
                    elif export_type == "segmentation" and ann['mask'] is not None:
                        # Format: class_id polygon_points...
                        mask = ann['mask']
                        seg_coords = mask_to_yolo_segmentation(mask, img_width, img_height)
                        if seg_coords:
                            line = f"0 {' '.join(map(str, seg_coords))}\n"
                            f.write(line)
                            exported_count += 1
            
            print(f"        ✅ Created: {label_file}")
        except Exception as e:
            print(f"        ❌ Error creating {label_file}: {e}")
            raise
    
    print(f"      ✅ YOLO export completed. Exported {exported_count} annotations")


# =============================================================
# HTML Template
# =============================================================

HTML_TEMPLATE = r"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Grounded SAM 2 - Batch Mode</title>
    <style>
        :root { --bg: #0f172a; --card: #1e293b; --text: #e2e8f0; --accent: #38bdf8; }
        body { font-family: system-ui, sans-serif; margin: 0; padding: 20px; background: var(--bg); color: var(--text); }
        .container { max-width: 1400px; margin: 0 auto; }
        
        .header { background: var(--card); padding: 20px; border-radius: 12px; margin-bottom: 24px; border: 1px solid #334155; }
        .section-title { font-size: 0.9rem; font-weight: 600; color: #94a3b8; margin-top: 16px; margin-bottom: 12px; }
        
        .controls { display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 16px; align-items: end; }
        
        label { display: block; font-size: 0.85rem; color: #94a3b8; margin-bottom: 6px; font-weight: 600; }
        input[type="text"], select {
            width: 100%; padding: 10px; border-radius: 8px; border: 1px solid #475569;
            background: #020617; color: white; font-family: monospace; font-size: 0.9rem;
        }
        select { font-family: system-ui; }
        
        .btn {
            padding: 10px 20px; border-radius: 8px; border: none; font-weight: bold; cursor: pointer;
            transition: opacity 0.2s; color: #0f172a; white-space: nowrap;
        }
        .btn-load { background: #94a3b8; }
        .btn-run { background: linear-gradient(135deg, #38bdf8, #22c55e); width: 100%; }
        .btn-export { background: linear-gradient(135deg, #f59e0b, #d97706); width: 100%; }
        .btn:hover { opacity: 0.9; }
        .btn:disabled { opacity: 0.5; cursor: not-allowed; }

        .export-section { background: #1e293b; border: 1px solid #334155; border-radius: 8px; padding: 16px; margin-top: 20px; }
        
        .gallery {
            display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 20px;
        }
        .img-card {
            background: var(--card); border-radius: 12px; overflow: hidden;
            border: 1px solid #334155; display: flex; flex-direction: column;
        }
        .img-header { padding: 10px; font-size: 0.8rem; color: #94a3b8; border-bottom: 1px solid #334155; display: flex; justify-content: space-between; }
        .img-container { position: relative; aspect-ratio: 4/3; background: #020617; }
        .img-container img { width: 100%; height: 100%; object-fit: contain; }
        .tag { background: #22c55e; color: black; padding: 2px 8px; border-radius: 99px; font-size: 0.7rem; font-weight: bold; }
        
        .empty-state { text-align: center; padding: 60px; color: #64748b; }
        .success-msg { background: #22c55e30; border: 1px solid #22c55e; color: #86efac; padding: 15px; border-radius: 8px; margin-bottom: 20px; }
        .error-msg { background: #ef444420; border: 1px solid #ef4444; color: #fca5a5; padding: 15px; border-radius: 8px; margin-bottom: 20px; }
    </style>
</head>
<body>
<div class="container">
    <div class="header">
        <h1>Grounded SAM 2 <span style="color:var(--accent)">Batch Processor</span></h1>
        
        <form method="GET" action="/">
            <div class="section-title">Input & Detection</div>
            <div class="controls">
                <div>
                    <label>Folder Path (Local)</label>
                    <div style="display:flex; gap:10px;">
                        <input type="text" name="folder" value="{{ folder }}" placeholder="C:\Users\...\images" required>
                        <button type="submit" class="btn btn-load">Load</button>
                    </div>
                </div>
                <div>
                    <label>Detection Prompt</label>
                    <input type="text" name="prompt" value="{{ prompt }}" placeholder="e.g. cat, dog, defect">
                </div>
                <div>
                    <button type="submit" name="action" value="run" class="btn btn-run">⚡ Run Inference</button>
                </div>
            </div>

            {% if folder and is_annotated %}
            <div class="export-section">
                <div class="section-title">Export Annotations</div>
                <div class="controls">
                    <div>
                        <label>Export Type</label>
                        <select name="export_type" required>
                            <option value="">-- Select --</option>
                            <option value="detection">Detection (Bounding Boxes)</option>
                            <option value="segmentation">Segmentation (Masks)</option>
                        </select>
                    </div>
                    <div>
                        <label>Label Format</label>
                        <select name="export_format" required>
                            <option value="">-- Select --</option>
                            <option value="coco">COCO (JSON)</option>
                            <option value="yolo">YOLO (TXT)</option>
                        </select>
                    </div>
                    <div>
                        <button type="submit" name="action" value="export" class="btn btn-export">📥 Export Labels</button>
                    </div>
                </div>
            </div>
            {% endif %}
        </form>
    </div>

    {% if success_msg %}
    <div class="success-msg">{{ success_msg }}</div>
    {% endif %}

    {% if error %}
    <div class="error-msg">{{ error }}</div>
    {% endif %}

    <div class="gallery">
        {% for img in images %}
        <div class="img-card">
            <div class="img-header">
                <span>{{ img.filename }}</span>
                {% if img.is_annotated %}
                <span class="tag">{{ img.annotations|length }} Detected</span>
                {% endif %}
            </div>
            <div class="img-container">
                <img src="data:image/jpeg;base64,{{ img.image_b64 }}" loading="lazy">
            </div>
        </div>
        {% else %}
            {% if folder %}
            <div class="empty-state" style="grid-column: 1/-1;">No images found in this folder.</div>
            {% endif %}
        {% endfor %}
    </div>

    {% if not folder %}
    <div class="empty-state">Enter a folder path above to begin.</div>
    {% endif %}
</div>
</body>
</html>
"""

# =============================================================
# Helper Functions
# =============================================================

def process_images_with_inference(folder: str, prompt: str) -> List[ProcessedImage]:
    """Process all images in folder with inference"""
    processed_images: List[ProcessedImage] = []
    filenames = sorted([
        f for f in os.listdir(folder) 
        if os.path.splitext(f)[1].lower() in ALLOWED_EXTENSIONS
    ])
    
    print(f"   Processing {len(filenames)} images with prompt: '{prompt}'")
    
    for idx, filename in enumerate(filenames):
        filepath = os.path.join(folder, filename)
        try:
            img = Image.open(filepath)
            print(f"   [{idx+1}/{len(filenames)}] {filename}...", end=" ", flush=True)
            
            res = run_grounded_dino_sam(img, prompt)
            final_img = res["annotated_image"]
            anns = res["annotations"]
            img_size = res["image_size"]
            
            print(f"({len(anns)} objects detected)")
            
            b64 = pil_to_base64(final_img)
            
            processed_images.append(ProcessedImage(
                filename=filename,
                image_b64=b64,
                annotations=anns,
                is_annotated=len(anns) > 0,
                image_size=img_size
            ))
        except Exception as e:
            print(f"ERROR: {e}")
            import traceback
            traceback.print_exc()
            continue
    
    return processed_images

# =============================================================
# Routes
# =============================================================

app = Flask(__name__)


@app.route("/", methods=["GET"])
def index():
    folder = request.args.get("folder", "").strip()
    prompt = request.args.get("prompt", "").strip()
    action = request.args.get("action", "")
    export_format = request.args.get("export_format", "").strip()
    export_type = request.args.get("export_type", "").strip()
    
    processed_images: List[ProcessedImage] = []
    error_msg = ""
    success_msg = ""
    has_any_annotations = False

    if folder:
        if not os.path.exists(folder):
            error_msg = f"Folder not found: {folder}"
        else:
            try:
                # 1. Scan folder for images
                filenames = sorted([
                    f for f in os.listdir(folder) 
                    if os.path.splitext(f)[1].lower() in ALLOWED_EXTENSIONS
                ])

                # 2. Process Loop
                for filename in filenames:
                    filepath = os.path.join(folder, filename)
                    
                    try:
                        img = Image.open(filepath)
                        
                        # Decide: Inference or Just Display?
                        if action == "run" and prompt:
                            # Run AI
                            res = run_grounded_dino_sam(img, prompt)
                            final_img = res["annotated_image"]
                            anns = res["annotations"]
                            img_size = res["image_size"]
                        else:
                            # Just Load
                            final_img = img
                            anns = []
                            img_size = img.size

                        # Track if any image has annotations
                        if anns:
                            has_any_annotations = True

                        # Convert to B64 for HTML
                        b64 = pil_to_base64(final_img)
                        
                        processed_images.append(ProcessedImage(
                            filename=filename,
                            image_b64=b64,
                            annotations=anns,
                            is_annotated=len(anns) > 0,
                            image_size=img_size
                        ))
                    except Exception as e:
                        print(f"Error reading {filename}: {e}")
                        continue

                # 3. Handle Export
                if action == "export" and export_format and export_type and prompt:
                    print(f"\n🔍 Export triggered: format={export_format}, type={export_type}, prompt={prompt}")
                    
                    # Re-run inference to get fresh annotations
                    print("   Re-running inference for export...")
                    processed_images = process_images_with_inference(folder, prompt)
                    
                    annotated_images = [img for img in processed_images if img.is_annotated and img.annotations]
                    print(f"   Annotated images found: {len(annotated_images)}/{len(processed_images)}")
                    
                    if not annotated_images:
                        error_msg = "❌ No objects detected. Try a different prompt or check your images."
                    else:
                        try:
                            # Convert ProcessedImage to dict for export
                            images_data = [img.to_export_dict() for img in processed_images]
                            print(f"   Images data prepared: {len(images_data)} images")
                            
                            export_labels(folder, images_data, export_format, export_type, prompt)
                            
                            labels_folder = os.path.join(os.path.dirname(folder), "labels")
                            if export_format == "coco":
                                success_msg = f"✅ Labels exported to {labels_folder}/annotations.json"
                            else:
                                success_msg = f"✅ Labels exported to {labels_folder}/ (one .txt per image)"
                            print(f"   ✅ Export successful")
                        except Exception as e:
                            print(f"   ❌ Export error: {str(e)}")
                            import traceback
                            traceback.print_exc()
                            error_msg = f"Error exporting labels: {str(e)}"

            except Exception as e:
                error_msg = f"Error scanning folder: {str(e)}"

    return render_template_string(
        HTML_TEMPLATE,
        folder=folder,
        prompt=prompt,
        images=processed_images,
        error=error_msg,
        success_msg=success_msg,
        is_annotated=has_any_annotations
    )

if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5000, debug=False)