model_interface/a_17_grapes_count.py

import streamlit as st
import os
import cv2
import pandas as pd
from ultralytics import YOLO
import numpy as np
from fpdf import FPDF
from model_interface.hf_model_store import get_artifact_path
import base64
import math
import tempfile

# =====================================================================
# CONFIG — resolve each image path independently via get_artifact_path
# =====================================================================
MODEL_PATH = get_artifact_path(r"17_grapes_count/best (2).pt")

IMAGES_DATA = {}
for img_num in [1, 2]:
    key = f"Image {img_num}"
    try:
        path = get_artifact_path(f"17_grapes_count/data/image_{img_num}.txt")
        if path and os.path.exists(path):
            IMAGES_DATA[key] = path
            print(f"[OK] {key} -> {path}")
        else:
            print(f"[MISSING] {key} resolved to path={path} but file not found")
    except Exception as e:
        print(f"[ERROR] {key}: {e}")

print(f"Available images: {list(IMAGES_DATA.keys())}")

# =====================================================================
# GRADING TABLES
# =====================================================================
COMMON_GRADING_SCALE = [
    {"label": "Fail",        "min": 0,  "max": 30},
    {"label": "On Hold",     "min": 31, "max": 40},
    {"label": "Poor",        "min": 41, "max": 55},
    {"label": "Acceptable",  "min": 56, "max": 72},
    {"label": "Fairly Good", "min": 73, "max": 85},
    {"label": "Good",        "min": 86, "max": 90},
    {"label": "Excellent",   "min": 91, "max": 100},
]

QC_TEMPLATES = {
    "Elongated White Seedless 2026": {
        "grading_scale": COMMON_GRADING_SCALE,
        "parameters": {
            "Berry Color": [(0, 1.9, 0), (2, 2.4, 5), (2.41, 3, 10)],
            "Berry Size":  [(0, 12.99, 0), (13, 25, 5), (25.01, 100, 0)],
        },
    },
    "Red Seedless 2026": {
        "grading_scale": COMMON_GRADING_SCALE,
        "parameters": {
            "Berry Color": [(0, 0.9, 0), (1, 1.6, 5), (1.61, 2, 10)],
            "Berry Size":  [(0, 15.99, 0), (16, 25, 10), (25.01, 100, 0)],
        },
    },
    "White Seedless 2026": {
        "grading_scale": COMMON_GRADING_SCALE,
        "parameters": {
            "Berry Color": [(0, 1.99, 0), (2, 2.6, 5), (2.61, 3, 10)],
            "Berry Size":  [(0, 14.9, 0), (15, 25, 10), (25.01, 100, 0)],
        },
    },
    "Black Seedless 2026": {
        "grading_scale": COMMON_GRADING_SCALE,
        "parameters": {
            "Berry Color": [(0, 1.39, 0), (1.4, 1.7, 5), (1.71, 2, 10)],
            "Berry Size":  [(0, 15.99, 0), (16, 25, 10), (25.01, 100, 0)],
        },
    },
    "Elongated White Seedless 2025": {
        "grading_scale": COMMON_GRADING_SCALE,
        "parameters": {
            "Berry Color": [(0, 1.9, 0), (2, 2.4, 5), (2.41, 3, 10)],
            "Berry Size":  [(0, 12.99, 0), (13, 25, 5), (25.01, 100, 0)],
        },
    },
}

# =====================================================================
# HELPERS
# =====================================================================

def load_base64_image_from_file(txt_file_path: str):
    """Read a base64-encoded .txt file and return decoded image bytes."""
    try:
        with open(txt_file_path, "rb") as f:
            raw = f.read().strip()
        return base64.b64decode(raw)
    except Exception as e:
        st.error(f"Could not decode image from {txt_file_path}: {e}")
        return None


def bytes_to_temp_jpg(image_bytes: bytes):
    """Write image bytes to a temp .jpg file and return its path."""
    try:
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
        tmp.write(image_bytes)
        tmp.close()
        return tmp.name
    except Exception as e:
        st.error(f"Error writing temp image: {e}")
        return None


# =====================================================================
# DETECTION PIPELINE
# =====================================================================

def grape_full_pipeline(model_path, image_path,
                        output_image="result.jpg",
                        excel_output="result.xlsx"):
    try:
        model = YOLO(model_path)
        img   = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"cv2.imread returned None for: {image_path}")

        results = model(img, conf=0.50, iou=0.5, imgsz=1280, max_det=1500)
        if not results or results[0].boxes is None:
            raise ValueError("No detections returned by model.")

        boxes       = results[0].boxes
        areas       = []
        green_count = 0
        black_count = 0
        excel_data  = []
        grape_id    = 0

        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            if x2 <= x1 or y2 <= y1:
                continue
            area = (x2 - x1) * (y2 - y1)
            if area < 150:
                continue
            crop = img[y1:y2, x1:x2]
            if crop.size == 0:
                continue

            hsv      = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV).reshape(-1, 3)
            total_px = len(hsv)
            if total_px == 0:
                continue

            h, s = hsv[:, 0], hsv[:, 1]
            green_px  = int(np.sum((h >= 20) & (h <  90) & (s > 40)))
            blue_px   = int(np.sum((h >= 90) & (h < 130)))
            purple_px = int(np.sum((h >= 130) & (h < 160)))
            red_px    = int(np.sum((h <  20) | (h >= 160)))
            other_px  = max(0, total_px - (green_px + blue_px + purple_px + red_px))

            areas.append(area)
            grape_id += 1

            if (green_px / total_px) * 100 > 20:
                color      = (0, 255, 0)
                label      = f"Green {grape_id}"
                grape_type = "Green"
                green_count += 1
            else:
                color      = (255, 0, 255)
                label      = f"Black {grape_id}"
                grape_type = "Black"
                black_count += 1

            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv2.putText(img, label, (x1, max(20, y1 - 5)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            excel_data.append({
                "grape_id":      grape_id,
                "type":          grape_type,
                "area":          area,
                "total_pixels":  total_px,
                "green_pixels":  green_px,
                "red_pixels":    red_px,
                "blue_pixels":   blue_px,
                "purple_pixels": purple_px,
                "other_pixels":  other_px,
            })

        total_grapes = len(areas)
        avg_size     = float(np.mean(areas)) if areas else 0.0
        small = medium = big = 0
        for a in areas:
            if   a < 0.75 * avg_size: small  += 1
            elif a > 1.25 * avg_size: big    += 1
            else:                     medium += 1

        cv2.putText(img, f"Total: {total_grapes}",
                    (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        cv2.imwrite(output_image, img)

        df = pd.DataFrame(excel_data)
        df.to_excel(excel_output, index=False)

        return {
            "image": output_image,
            "excel": excel_output,
            "df":    df,
            "summary": {
                "total":    total_grapes,
                "green":    green_count,
                "black":    black_count,
                "avg_size": round(avg_size, 2),
                "small":    small,
                "medium":   medium,
                "big":      big,
            },
        }
    except Exception as e:
        return {"error": str(e)}


# =====================================================================
# GRADING
# =====================================================================

def calculate_grade(template_name, user_inputs):
    template      = QC_TEMPLATES[template_name]
    total_penalty = 0
    breakdown     = []

    for param, value in user_inputs.items():
        if param not in template["parameters"]:
            continue
        rules   = template["parameters"][param]
        penalty = 0
        if isinstance(rules, list):
            try:
                val = float(value)
                for (lo, hi, p) in rules:
                    if lo <= val <= hi:
                        penalty = p
                        break
            except (ValueError, TypeError):
                pass
        elif isinstance(rules, dict):
            penalty = rules.get(value, 0)

        total_penalty += penalty
        breakdown.append({"Parameter": param, "Input": value, "Penalty": penalty})

    score       = 100 - total_penalty
    grade_label = "Fail"
    for scale in template["grading_scale"]:
        if scale["min"] <= score <= scale["max"]:
            grade_label = scale["label"]
            break

    return score, grade_label, breakdown


# =====================================================================
# PDF REPORT
# =====================================================================

def generate_pdf_report(image_path, result, score, grade_label, breakdown, template_name):
    pdf = FPDF()
    pdf.add_page()

    pdf.set_font("Arial", "B", 16)
    pdf.cell(0, 10, "Grape Quality Control Report", ln=1, align="C")

    pdf.set_font("Arial", size=12)
    pdf.cell(0, 10, f"Template: {template_name}", ln=1)
    pdf.cell(0, 10, f"Final Score: {score} / 100",  ln=1)
    pdf.cell(0, 10, f"Grade: {grade_label}",         ln=1)
    pdf.ln(5)

    if os.path.exists(result["image"]):
        pdf.image(result["image"], w=140)
    pdf.ln(10)

    pdf.set_font("Arial", "B", 14)
    pdf.cell(0, 10, "Detection Summary", ln=1)
    pdf.set_font("Arial", size=12)
    s = result["summary"]
    for line in [f"Total Grapes : {s['total']}",
                 f"Green Grapes : {s['green']}",
                 f"Black Grapes : {s['black']}",
                 f"Average Size : {s['avg_size']}"]:
        pdf.cell(0, 8, line, ln=1)
    pdf.ln(10)

    pdf.set_font("Arial", "B", 14)
    pdf.cell(0, 10, "QC Grading Breakdown", ln=1)
    pdf.set_font("Arial", "B", 12)
    for header, w in [("Parameter", 60), ("Value", 60), ("Penalty", 40)]:
        pdf.cell(w, 8, header, 1)
    pdf.ln()
    pdf.set_font("Arial", size=12)
    for b in breakdown:
        pdf.cell(60, 8, str(b["Parameter"])[:28], 1)
        pdf.cell(60, 8, str(b["Input"])[:28],     1)
        pdf.cell(40, 8, str(b["Penalty"]),         1)
        pdf.ln()

    pdf_path = "report.pdf"
    pdf.output(pdf_path)
    return pdf_path


# =====================================================================
# STREAMLIT APP
# =====================================================================

def grapes_count():
    st.set_page_config(layout="wide", page_title="Grape QC Dashboard", page_icon="🍇")
    st.title("🍇 Grape Detection & QC Dashboard")

    # ── initialise session state ───────────────────────────────────────
    defaults = {
        "result":         None,
        "score":          None,
        "grade_label":    None,
        "breakdown":      None,
        "g_ratio":        0.0,
        "b_ratio":        0.0,
        "last_image":     None,
        "last_template":  None,
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v

    # ── sidebar ────────────────────────────────────────────────────────
    with st.sidebar:
        st.header("1. Image Selection")

        if IMAGES_DATA:
            # IMAGES_DATA keys are exactly ["Image 1", "Image 2"] for whichever exist
            image_options  = list(IMAGES_DATA.keys())
            selected_image = st.selectbox("Select Image", image_options, key="img_select")
        else:
            st.warning("⚠️ No sample images found.")
            st.code(
                "17_grapes_count/data/image_1.txt\n"
                "17_grapes_count/data/image_2.txt",
                language="text",
            )
            selected_image = None

        st.header("2. QC Template")
        selected_template = st.selectbox(
            "Choose QC Template", list(QC_TEMPLATES.keys()), key="tmpl_select"
        )

        run_btn = st.button(
            "▶ Run Prediction & Grade",
            type="primary",
            use_container_width=True,
        )

    # ── clear cached result when user changes image or template ────────
    selection_changed = (
        selected_image    != st.session_state["last_image"] or
        selected_template != st.session_state["last_template"]
    )
    if selection_changed:
        for k in ("result", "score", "grade_label", "breakdown"):
            st.session_state[k] = None
        st.session_state["g_ratio"]       = 0.0
        st.session_state["b_ratio"]       = 0.0
        st.session_state["last_image"]    = selected_image
        st.session_state["last_template"] = selected_template

    # ── load input image bytes → temp file ────────────────────────────
    image_path = None
    if selected_image and selected_image in IMAGES_DATA:
        img_bytes = load_base64_image_from_file(IMAGES_DATA[selected_image])
        if img_bytes:
            image_path = bytes_to_temp_jpg(img_bytes)

    if image_path is None:
        st.info("Select an image from the sidebar to begin.")
        return

    # ── run pipeline when button is pressed ───────────────────────────
    if run_btn:
        with st.spinner("Running YOLO detection…"):
            res = grape_full_pipeline(MODEL_PATH, image_path)

        if "error" in res:
            st.error(f"Detection failed: {res['error']}")
        else:
            avg_sz = res["summary"]["avg_size"]
            tot    = res["summary"]["total"]

            user_inputs = {
                "Berry Size": round(math.sqrt(avg_sz), 2) if avg_sz > 0 else 0.0,
            }
            if tot > 0:
                g_ratio = res["summary"]["green"] / tot
                b_ratio = res["summary"]["black"] / tot
            else:
                g_ratio = b_ratio = 0.0
            user_inputs["Berry Color"] = round(1.0 * g_ratio + 3.0 * b_ratio, 2)

            score, grade_label, breakdown = calculate_grade(selected_template, user_inputs)

            # Variety mismatch hard-fail
            tmpl_lower = selected_template.lower()
            if ("white" in tmpl_lower or "elongated" in tmpl_lower) and b_ratio > 0.50:
                score, grade_label = 0, "Fail - Wrong Variety"
                breakdown.append({"Parameter": "Variety Verification",
                                   "Input": "Red/Black Detected", "Penalty": 100})
            elif ("red" in tmpl_lower or "black" in tmpl_lower) and g_ratio > 0.50:
                score, grade_label = 0, "Fail - Wrong Variety"
                breakdown.append({"Parameter": "Variety Verification",
                                   "Input": "Green Detected", "Penalty": 100})

            st.session_state.update({
                "result":      res,
                "score":       score,
                "grade_label": grade_label,
                "breakdown":   breakdown,
                "g_ratio":     g_ratio,
                "b_ratio":     b_ratio,
            })

    # ── Read input image aspect ratio once ────────────────────────────
    _probe = cv2.imread(image_path)
    _aspect_pct = round((_probe.shape[0] / _probe.shape[1]) * 100, 2) \
                  if _probe is not None else 75.0

    # ── Tabs: Detection | Results — both always exist in the DOM ──────
    # Using tabs means the results panel is NEVER below the images, so
    # it can never push the image columns up and cause shaking.
    tab_detect, tab_results = st.tabs(["🔍 Detection", "📊 Results & Export"])

    # ── TAB 1: Detection ──────────────────────────────────────────────
    with tab_detect:
        col_input, col_output = st.columns(2)

        with col_input:
            st.subheader(f"📷 Input — {selected_image}")
            st.image(image_path, use_container_width=True)

        with col_output:
            st.subheader("🔍 Predicted Output")
            out_slot = st.empty()
            if st.session_state["result"] is not None:
                out_img = st.session_state["result"]["image"]
                if os.path.exists(out_img):
                    out_slot.image(out_img, use_container_width=True)
            else:
                out_slot.markdown(
                    f"""
                    <div style="position:relative;width:100%;
                                padding-top:{_aspect_pct}%;
                                background:#1e1e1e;border-radius:8px;
                                border:1px dashed #444;">
                      <div style="position:absolute;inset:0;display:flex;
                                  align-items:center;justify-content:center;
                                  color:#888;font-size:15px;">
                        Click <strong>▶ Run</strong> in the sidebar
                      </div>
                    </div>
                    """,
                    unsafe_allow_html=True,
                )

        # Status bar at bottom of detection tab — never causes layout shift
        if st.session_state["result"] is not None:
            st.success(f"✅ Detection complete — "
                       f"{st.session_state['result']['summary']['total']} grapes found. "
                       f"Switch to the **Results & Export** tab for full analysis.")
        else:
            st.info("Press **▶ Run Prediction & Grade** in the sidebar to start.")

    # ── TAB 2: Results & Export ───────────────────────────────────────
    with tab_results:
        if st.session_state["result"] is None:
            st.info("Run a prediction first to see results here.")
        else:
            res         = st.session_state["result"]
            score       = st.session_state["score"]
            grade_label = st.session_state["grade_label"]
            breakdown   = st.session_state["breakdown"]

            st.header("📊 Evaluation Results")
            m1, m2, m3, m4 = st.columns(4)
            m1.metric("Total Grapes",  res["summary"]["total"])
            m2.metric("Green / Black",
                      f'{res["summary"]["green"]} / {res["summary"]["black"]}')
            m3.metric("QC Score",      f"{score} / 100")
            m4.metric("Grade",         grade_label)

            st.markdown("---")
            t1, t2 = st.columns(2)
            with t1:
                st.subheader("QC Penalty Breakdown")
                st.dataframe(pd.DataFrame(breakdown), use_container_width=True)
            with t2:
                st.subheader("Detection Summary")
                st.json(res["summary"])

            st.subheader("📄 Raw Detection Data")
            st.dataframe(res["df"], use_container_width=True)

            st.markdown("---")
            st.header("📥 Export Reports")
            pdf_path = generate_pdf_report(
                image_path, res, score, grade_label, breakdown, selected_template
            )

            dl1, dl2 = st.columns(2)
            with dl1:
                with open(res["excel"], "rb") as f:
                    st.download_button(
                        label="⬇ Download Excel Data",
                        data=f,
                        file_name="grape_analysis.xlsx",
                        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                        use_container_width=True,
                    )
            with dl2:
                with open(pdf_path, "rb") as f:
                    st.download_button(
                        label="📄 Download PDF Report",
                        data=f,
                        file_name=f"QC_Report_{selected_template}.pdf",
                        mime="application/pdf",
                        use_container_width=True,
                    )


if __name__ == "__main__":
    grapes_count()