app.py

# app.py - COMPLETE FIXED VERSION

import json
import logging
import base64
import io
import torch
import gc
from pathlib import Path
from typing import List, Dict, Any

from threading import Lock, Thread

import numpy as np
from flask import Flask, render_template, request, jsonify
from PIL import Image, ImageDraw

from config import BASE_DIR
from exporters import export_labels
from inference import (
    load_images_without_inference,
    process_images_with_inference,
    run_grounded_dino_sam,
    sam2_predictor,
    set_active_model_path,
    unload_models,
    SAM_DTYPE,
    DEVICE,
)
from models import ProcessedImage
from utils import pil_to_base64
from train_dino import FineTuner

# Basic logging configuration
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)

logger = logging.getLogger(__name__)

# Flask application
app = Flask(__name__, template_folder=str(BASE_DIR / "templates"))

SESSION_FILE_NAME = "_annotations_session.json"
ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

# Global fine-tune training state
training_state: Dict[str, Any] = {
    "running": False,
    "status": "idle",
    "current_epoch": 0,
    "total_epochs": 0,
    "loss": None,
    "last_checkpoint": None,
    "error": None,
}
training_lock = Lock()


# ---------------------------------------------------------------------
# Session helpers
# ---------------------------------------------------------------------
def get_session_path(folder: Path) -> Path:
    return folder / SESSION_FILE_NAME


def load_session_annotations(folder: Path) -> List[Dict]:
    session_path = get_session_path(folder)
    if not session_path.is_file():
        return []
    try:
        with session_path.open("r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as exc:
        logger.exception("Failed to read session file %s: %s", session_path, exc)
        return []


def save_session_annotations(folder: Path, images_data: List[Dict]) -> None:
    session_path = get_session_path(folder)
    try:
        with session_path.open("w", encoding="utf-8") as f:
            json.dump(images_data, f, indent=2)
    except Exception as exc:
        logger.exception("Failed to write session file %s: %s", session_path, exc)


def processed_images_to_session_data(
    processed_images: List[ProcessedImage],
) -> List[Dict[str, Any]]:
    session_data: List[Dict[str, Any]] = []
    for img in processed_images:
        w, h = img.image_size
        anns_list = []
        if img.annotations:
            for ann in img.annotations:
                ann_dict = {
                    "label": getattr(ann, "label", ""),
                    "score": float(getattr(ann, "score", 1.0)),
                    "bbox": [float(x) for x in getattr(ann, "bbox", [0.0, 0.0, 0.0, 0.0])],
                }
                if hasattr(ann, "mask") and ann.mask is not None:
                    try:
                        mask_u8 = (ann.mask.astype(np.uint8) * 255)
                        mask_pil = Image.fromarray(mask_u8)
                        buf = io.BytesIO()
                        mask_pil.save(buf, format="PNG")
                        mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                        ann_dict["mask_b64"] = mask_b64
                    except Exception as e:
                        logger.warning(f"Failed to encode mask for {img.filename}: {e}")
                anns_list.append(ann_dict)

        session_data.append(
            {
                "filename": img.filename,
                "image_size": [int(w), int(h)],
                "is_annotated": bool(img.is_annotated),
                "image_b64": img.image_b64,
                "annotations": anns_list,
            }
        )
    return session_data


def list_image_files(folder: Path) -> List[str]:
    return sorted(
        p.name
        for p in folder.iterdir()
        if p.is_file() and p.suffix.lower() in ALLOWED_EXTENSIONS
    )


def session_to_processed_images(
    folder: Path, session_images_data: List[Dict]
) -> List[ProcessedImage]:
    processed: List[ProcessedImage] = []

    for item in session_images_data:
        filename = item.get("filename", "")
        image_b64 = item.get("image_b64", "")
        img_size = item.get("image_size") or [0, 0]
        is_annotated = bool(item.get("is_annotated"))
        annotations = item.get("annotations") or []

        if not image_b64 and filename:
            img_path = folder / filename
            try:
                img = Image.open(img_path).convert("RGB")
                image_b64 = pil_to_base64(img)
            except Exception as exc:
                image_b64 = ""

        if len(img_size) != 2:
            img_size = [0, 0]

        processed.append(
            ProcessedImage(
                filename=filename,
                image_b64=image_b64,
                annotations=annotations,
                is_annotated=is_annotated,
                image_size=tuple(img_size),
            )
        )

    return processed


def render_image_with_annotations(folder: Path, filename: str, annotations: List[Dict]) -> str:
    img_path = folder / filename
    try:
        img = Image.open(img_path).convert("RGB")
    except Exception:
        return ""

    mask_layer = Image.new("RGBA", img.size, (0, 0, 0, 0))
    draw_img = ImageDraw.Draw(img)
    
    for ann in annotations:
        mask_b64 = ann.get("mask_b64")
        if mask_b64:
            try:
                mask_data = base64.b64decode(mask_b64)
                mask_pil = Image.open(io.BytesIO(mask_data))
                if mask_pil.size != img.size:
                    mask_pil = mask_pil.resize(img.size, resample=Image.NEAREST)
                
                solid_color = Image.new("RGBA", img.size, (56, 189, 248, 100))
                mask_layer.paste(solid_color, (0, 0), mask_pil)
            except Exception as e:
                logger.warning(f"Failed to decode mask: {e}")

        bbox = ann.get("bbox") or [0, 0, 0, 0]
        x1, y1, x2, y2 = bbox
        draw_img.rectangle([x1, y1, x2, y2], outline=(56, 189, 248), width=3)
        
        label = ann.get("label", "")
        draw_img.text((x1 + 2, y1 - 18), label, fill=(0, 0, 0))

    img = img.convert("RGBA")
    img = Image.alpha_composite(img, mask_layer)
    img = img.convert("RGB")

    return pil_to_base64(img)


# ---------------------------------------------------------------------
# Main Routes
# ---------------------------------------------------------------------
@app.route("/", methods=["GET"])
def index():
    folder = request.args.get("folder", "").strip()
    prompt = request.args.get("prompt", "").strip()

    folder_path = Path(folder)
    if folder and not folder_path.is_dir():
        return render_template("index.html", error="Invalid folder path", images=[])

    session_images_data = load_session_annotations(folder_path) if folder else []
    images = []
    if session_images_data:
        images = session_to_processed_images(folder_path, session_images_data)
    elif folder:
        images = load_images_without_inference(folder)

    return render_template(
        "index.html",
        folder=folder,
        prompt=prompt,
        images=images,
        is_annotated=any(img.is_annotated for img in images) if images else False,
    )


@app.route("/edit/<path:folder>/<filename>")
def edit(folder, filename):
    folder_path = Path(folder)
    prompt = request.args.get("prompt", "")

    if not folder_path.is_dir():
        return "Invalid folder", 400

    img_path = folder_path / filename
    if not img_path.is_file():
        return "Image not found", 404

    session_images_data = load_session_annotations(folder_path)
    annotations = []
    for item in session_images_data:
        if item.get("filename") == filename:
            annotations = item.get("annotations") or []
            break
            
    try:
        img = Image.open(img_path).convert("RGB")
        w, h = img.size
        image_b64 = pil_to_base64(img)
    except Exception as e:
        return f"Error loading image: {e}", 500

    return render_template(
        "edit.html",
        folder=folder,
        filename=filename,
        prompt=prompt,
        image_b64=image_b64,
        image_width=w,
        image_height=h,
        annotations_json=json.dumps(annotations),
    )


# ---------------------------------------------------------------------
# API Routes
# ---------------------------------------------------------------------
@app.route("/api/load_images", methods=["POST"])
def api_load_images():
    data = request.get_json(force=True)
    folder = data.get("folder", "").strip()
    if not folder: 
        return jsonify({"status": "error", "message": "No folder provided"}), 400
    folder_path = Path(folder)
    if not folder_path.is_dir(): 
        return jsonify({"status": "error", "message": "Invalid folder path"}), 400

    images = load_images_without_inference(str(folder_path))
    session_images_data = processed_images_to_session_data(images)
    save_session_annotations(folder_path, session_images_data)

    return jsonify({"status": "ok", "images": [img.__dict__ for img in images]})


@app.route("/api/run_inference", methods=["POST"])
def api_run_inference():
    data = request.get_json(force=True)
    folder = data.get("folder", "").strip()
    prompt = data.get("prompt", "").strip()
    model_checkpoint = data.get("model_checkpoint")

    if not folder or not prompt: 
        return jsonify({"status": "error", "message": "Missing folder or prompt"}), 400

    processed_images = process_images_with_inference(str(folder), prompt, model_checkpoint)
    session_images_data = processed_images_to_session_data(processed_images)
    save_session_annotations(Path(folder), session_images_data)

    return jsonify({"status": "ok", "images": [img.__dict__ for img in processed_images]})


@app.route("/api/export_labels", methods=["POST"])
def api_export_labels():
    data = request.get_json(force=True)
    folder = data.get("folder", "").strip()
    export_format = data.get("format", "coco").lower()
    export_type = data.get("type", "detection").lower()
    prompt = data.get("prompt", "").strip()

    folder_path = Path(folder)
    session_images_data = load_session_annotations(folder_path)
    
    try:
        export_labels(str(folder), session_images_data, export_format, export_type, prompt)
        return jsonify({"status": "ok"})
    except Exception as e:
        logger.exception("Export failed")
        return jsonify({"status": "error", "message": str(e)}), 500


@app.route("/api/save_annotations", methods=["POST"])
def save_annotations():
    data = request.get_json(force=True)
    folder = (data.get("folder") or "").strip()
    filename = (data.get("filename") or "").strip()
    annotations = data.get("annotations", [])
    folder_path = Path(folder)

    img_path = folder_path / filename
    try:
        img = Image.open(img_path).convert("RGB")
        width, height = img.size
    except Exception as e:
        logger.error(f"Failed to open image: {e}")
        return jsonify({"status": "error", "message": "Failed to load image"}), 500

    if annotations:
        image_b64 = render_image_with_annotations(folder_path, filename, annotations)
    else:
        image_b64 = pil_to_base64(img)

    session_images_data = load_session_annotations(folder_path)
    updated = False
    for item in session_images_data:
        if item.get("filename") == filename:
            item["annotations"] = annotations
            item["is_annotated"] = bool(annotations)
            item["image_b64"] = image_b64
            item["image_size"] = [width, height]
            updated = True
            break

    if not updated:
        session_images_data.append({
            "filename": filename,
            "image_size": [width, height],
            "is_annotated": bool(annotations),
            "annotations": annotations,
            "image_b64": image_b64,
        })

    save_session_annotations(folder_path, session_images_data)
    return jsonify({"status": "ok", "image_b64": image_b64})


@app.route("/api/clear_annotations", methods=["POST"])
def clear_annotations():
    data = request.get_json(force=True)
    folder = (data.get("folder") or "").strip()
    filename = (data.get("filename") or "").strip()
    folder_path = Path(folder)
    img_path = folder_path / filename
    
    try:
        img = Image.open(img_path).convert("RGB")
        width, height = img.size
        image_b64 = pil_to_base64(img)
    except Exception as e:
        logger.error(f"Failed to open image: {e}")
        return jsonify({"status": "error", "message": "Failed to load image"}), 500

    session_images_data = load_session_annotations(folder_path)
    for item in session_images_data:
        if item.get("filename") == filename:
            item["annotations"] = []
            item["is_annotated"] = False
            item["image_b64"] = image_b64
            break
    save_session_annotations(folder_path, session_images_data)
    return jsonify({"status": "ok", "image_b64": image_b64})


# ---------------------------------------------------------------------
# SAM 2 Interactive Routes (FIXED with proper autocast)
# ---------------------------------------------------------------------
@app.route("/api/auto_segment", methods=["POST"])
def api_auto_segment():
    data = request.get_json(force=True)
    folder = data.get("folder", "").strip()
    filename = data.get("filename", "").strip()
    bbox = data.get("bbox", [])
    ann_idx = data.get("annotation_index", -1)

    folder_path = Path(folder)
    img_path = folder_path / filename
    
    if sam2_predictor is None:
        return jsonify({"status": "error", "message": "SAM2 model not loaded"}), 500

    try:
        img = Image.open(img_path).convert("RGB")
        img_np = np.array(img)
        x1, y1, x2, y2 = [float(v) for v in bbox]

        # FIXED: Use autocast only for CUDA
        with torch.inference_mode():
            if DEVICE == "cuda":
                with torch.autocast(device_type="cuda", dtype=SAM_DTYPE):
                    sam2_predictor.set_image(img_np)
                    masks, scores, _ = sam2_predictor.predict(
                        box=np.array([[x1, y1, x2, y2]]),
                        multimask_output=False,
                    )
            else:
                sam2_predictor.set_image(img_np)
                masks, scores, _ = sam2_predictor.predict(
                    box=np.array([[x1, y1, x2, y2]]),
                    multimask_output=False,
                )
        
        if masks.ndim == 4: 
            masks = masks.squeeze(1)
        if masks.shape[0] == 0:
            return jsonify({"status": "error", "message": "No mask returned"}), 500

        mask = masks[0] > 0.0
        
        # Encode mask
        mask_img = (mask.astype(np.uint8) * 255)
        mask_pil = Image.fromarray(mask_img)
        buf = io.BytesIO()
        mask_pil.save(buf, format="PNG")
        mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        # Update Session
        session_images_data = load_session_annotations(folder_path)
        for item in session_images_data:
            if item.get("filename") == filename:
                annotations = item.get("annotations") or []
                if 0 <= ann_idx < len(annotations):
                    annotations[ann_idx]["mask_b64"] = mask_b64
                item["image_b64"] = render_image_with_annotations(folder_path, filename, annotations)
                item["annotations"] = annotations
                break
        save_session_annotations(folder_path, session_images_data)
        
        updated_image_b64 = ""
        for item in session_images_data:
            if item.get("filename") == filename:
                updated_image_b64 = item.get("image_b64", "")
                break

        return jsonify({"status": "ok", "mask_b64": mask_b64, "image_b64": updated_image_b64})
    except Exception as e:
        logger.exception("Auto-segment failed")
        return jsonify({"status": "error", "message": str(e)}), 500


@app.route("/api/refine_mask", methods=["POST"])
def api_refine_mask():
    data = request.get_json(force=True)
    folder = data.get("folder", "").strip()
    filename = data.get("filename", "").strip()
    ann_idx = data.get("annotation_index", -1)
    
    positive_points = data.get("positive_points", [])
    negative_points = data.get("negative_points", [])

    folder_path = Path(folder)
    img_path = folder_path / filename
    
    if sam2_predictor is None:
        return jsonify({"status": "error", "message": "SAM2 model not loaded"}), 500

    try:
        session_images_data = load_session_annotations(folder_path)
        bbox = None
        for item in session_images_data:
            if item.get("filename") == filename:
                annotations = item.get("annotations") or []
                if 0 <= ann_idx < len(annotations):
                    bbox = annotations[ann_idx].get("bbox")
                break
        
        if not bbox: 
            return jsonify({"status": "error", "message": "No bbox found"}), 404

        img = Image.open(img_path).convert("RGB")
        img_np = np.array(img)
        x1, y1, x2, y2 = bbox
        
        point_coords = []
        point_labels = []
        for p in positive_points:
            point_coords.append([p['x'], p['y']])
            point_labels.append(1)
        for p in negative_points:
            point_coords.append([p['x'], p['y']])
            point_labels.append(0)

        if not point_coords: 
            return jsonify({"status": "error", "message": "No points"}), 400

        # FIXED: Use autocast only for CUDA
        with torch.inference_mode():
            if DEVICE == "cuda":
                with torch.autocast(device_type="cuda", dtype=SAM_DTYPE):
                    sam2_predictor.set_image(img_np)
                    masks, scores, _ = sam2_predictor.predict(
                        point_coords=np.array(point_coords),
                        point_labels=np.array(point_labels),
                        box=np.array([[x1, y1, x2, y2]]),
                        multimask_output=False,
                    )
            else:
                sam2_predictor.set_image(img_np)
                masks, scores, _ = sam2_predictor.predict(
                    point_coords=np.array(point_coords),
                    point_labels=np.array(point_labels),
                    box=np.array([[x1, y1, x2, y2]]),
                    multimask_output=False,
                )
        
        if masks.ndim == 4: 
            masks = masks.squeeze(1)
        mask = masks[0] > 0.0

        mask_img = (mask.astype(np.uint8) * 255)
        mask_pil = Image.fromarray(mask_img)
        buf = io.BytesIO()
        mask_pil.save(buf, format="PNG")
        mask_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
        
        return jsonify({"status": "ok", "mask_b64": mask_b64})

    except Exception as e:
        logger.exception("Refine mask failed")
        return jsonify({"status": "error", "message": str(e)}), 500


# ---------------------------------------------------------------------
# Fine-tune & Batch Processing
# ---------------------------------------------------------------------
@app.route("/api/fine_tune", methods=["POST"])
def api_fine_tune():
    data = request.get_json(force=True)
    folder = (data.get("folder") or "").strip()
    prompt = (data.get("prompt") or "").strip()
    epochs = int(data.get("epochs") or 5)
    batch_size = int(data.get("batch_size") or 2)
    folder_path = Path(folder)

    with training_lock:
        if training_state.get("running"):
            return jsonify({"status": "busy", "message": "Training already running"}), 400
        training_state.update({
            "running": True, 
            "status": "starting", 
            "error": None,
            "current_epoch": 0, 
            "total_epochs": epochs,
            "loss": None,
            "last_checkpoint": None,
        })

    t = Thread(
        target=_run_finetune_thread,
        args=(folder_path, prompt, epochs, batch_size),
        daemon=True,
    )
    t.start()
    return jsonify({"status": "ok"})


@app.route("/api/fine_tune_status", methods=["GET"])
def api_fine_tune_status():
    with training_lock:
        return jsonify(training_state)


@app.route("/api/load_checkpoint", methods=["POST"])
def api_load_checkpoint():
    data = request.get_json(force=True)
    checkpoint_path = data.get("checkpoint_path")
    set_active_model_path(checkpoint_path)
    return jsonify({"status": "ok"})


@app.route("/api/reset_model", methods=["POST"])
def api_reset_model():
    set_active_model_path(None)
    return jsonify({"status": "ok"})


def _run_finetune_thread(folder_path: Path, prompt: str, epochs: int, batch_size: int):
    """Background thread for fine-tuning."""
    try:
        with training_lock:
            unload_models()
    except Exception as e:
        logger.warning(f"Failed to unload models: {e}")

    output_dir = folder_path / "checkpoints"
    ft = FineTuner(str(folder_path), prompt, str(output_dir))

    def progress_cb(progress: Dict[str, Any]):
        with training_lock:
            phase = progress.get("phase")
            if phase == "running":
                training_state.update({
                    "status": "running", 
                    "current_epoch": progress.get("epoch"), 
                    "loss": progress.get("loss")
                })
            elif phase == "done":
                training_state.update({
                    "running": False, 
                    "status": "completed", 
                    "last_checkpoint": progress.get("checkpoint_path")
                })
            elif phase == "error":
                training_state.update({
                    "running": False, 
                    "status": "error", 
                    "error": progress.get("error")
                })

    try:
        ft.fine_tune(num_epochs=epochs, batch_size=batch_size, progress_callback=progress_cb)
    except Exception as e:
        logger.exception("Fine-tuning failed")
        with training_lock:
            training_state.update({
                "running": False, 
                "status": "error", 
                "error": str(e)
            })


@app.route("/api/batch_process", methods=["POST"])
def api_batch_process():
    """FIXED: Batch process all images with better memory management."""
    data = request.get_json(force=True)
    folder = (data.get("folder") or "").strip()
    prompt = (data.get("prompt") or "").strip()
    use_finetuned = data.get("use_finetuned", False)
    folder_path = Path(folder)

    if not folder or not prompt:
        return jsonify({"status": "error", "message": "Missing folder or prompt"}), 400

    try:
        ckpt = None
        if use_finetuned:
            ft_ckpt = folder_path / "checkpoints" / "groundingdino_finetuned.pth"
            if ft_ckpt.exists(): 
                ckpt = str(ft_ckpt)
                logger.info(f"Using fine-tuned checkpoint: {ckpt}")
        
        set_active_model_path(ckpt)

        session_images_data = load_session_annotations(folder_path)
        all_files = list_image_files(folder_path)
        session_map = {item["filename"]: i for i, item in enumerate(session_images_data)}
        
        # Clean VRAM before starting
        torch.cuda.empty_cache()
        gc.collect()
        
        processed_count = 0
        failed_images = []

        for idx, filename in enumerate(all_files, 1):
            try:
                logger.info(f"Processing {idx}/{len(all_files)}: {filename}")
                img = Image.open(folder_path / filename).convert("RGB")
                
                # Inference with proper error handling
                res = run_grounded_dino_sam(img, prompt)
                
                anns_list = []
                for ann in res["annotations"]:
                    ad = {
                        "label": ann.label,
                        "score": float(ann.score),
                        "bbox": [float(x) for x in ann.bbox],
                    }
                    if ann.mask is not None:
                        buf = io.BytesIO()
                        Image.fromarray((ann.mask.astype(np.uint8)*255)).save(buf, format="PNG")
                        ad["mask_b64"] = base64.b64encode(buf.getvalue()).decode("utf-8")
                    anns_list.append(ad)

                # Render image with annotations
                if anns_list:
                    image_b64 = render_image_with_annotations(folder_path, filename, anns_list)
                else:
                    image_b64 = pil_to_base64(img)

                new_entry = {
                    "filename": filename,
                    "image_size": list(img.size),
                    "is_annotated": bool(anns_list),
                    "annotations": anns_list,
                    "image_b64": image_b64,
                }

                if filename in session_map:
                    session_images_data[session_map[filename]] = new_entry
                else:
                    session_images_data.append(new_entry)
                    session_map[filename] = len(session_images_data) - 1

                processed_count += 1
                
                # Periodically save and clean memory
                if processed_count % 5 == 0:
                    save_session_annotations(folder_path, session_images_data)
                    torch.cuda.empty_cache()
                    gc.collect()
                    logger.info(f"✅ Checkpoint: Processed {processed_count}/{len(all_files)}")

            except torch.cuda.OutOfMemoryError:
                logger.error(f"❌ OOM on {filename}, attempting recovery...")
                failed_images.append(filename)
                torch.cuda.empty_cache()
                gc.collect()
                # Try to reload models after OOM
                try:
                    unload_models()
                    set_active_model_path(ckpt)
                except Exception as e:
                    logger.error(f"Failed to reload models: {e}")
            except Exception as e:
                logger.error(f"❌ Failed to process {filename}: {e}")
                failed_images.append(filename)
                torch.cuda.empty_cache()

        # Final save
        save_session_annotations(folder_path, session_images_data)
        
        # Reset to base model
        set_active_model_path(None)
        torch.cuda.empty_cache()

        message = f"Processed {processed_count}/{len(all_files)} images"
        if failed_images:
            message += f" ({len(failed_images)} failed)"
            logger.warning(f"Failed images: {failed_images}")

        return jsonify({
            "status": "ok",
            "message": message,
            "processed": processed_count,
            "total": len(all_files),
            "failed": len(failed_images),
            "failed_images": failed_images[:10]  # Only return first 10 to avoid large response
        })

    except Exception as e:
        logger.exception("Error in batch processing")
        return jsonify({"status": "error", "message": str(e)}), 500


@app.route("/api/get_labeling_stats", methods=["POST"])
def api_get_labeling_stats():
    data = request.get_json(force=True)
    folder = (data.get("folder") or "").strip()
    if not folder: 
        return jsonify({"status": "error", "message": "No folder provided"}), 400
    folder_path = Path(folder)
    
    try:
        session = load_session_annotations(folder_path)
        all_files = list_image_files(folder_path)
        labeled = len([i for i in session if i.get("is_annotated")])
        return jsonify({
            "status": "ok",
            "total_images": len(all_files),
            "labeled": labeled,
            "unlabeled": len(all_files) - labeled
        })
    except Exception as e:
        logger.exception("Failed to get labeling stats")
        return jsonify({"status": "error", "message": str(e)}), 500


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