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"""
Click β–Ά Run in the sidebar
""", 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()