Update app.py

app.py

@@ -1,16 +1,26 @@
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
 import os, cv2, re, base64
 import numpy as np
 import pandas as pd
 import gradio as gr
+from difflib import get_close_matches
 from roboflow import Roboflow
 from openai import OpenAI
 from openpyxl import load_workbook
+import torch.nn.functional as F
+from PIL import Image
+from concurrent.futures import ThreadPoolExecutor
+import cv2
 
 # ================= CONFIG =================
 
 ROBOFLOW_API_KEY = "uP19IAi98TqwLvHmNB8V"
 ROBOFLOW_PROJECT = "braker3"
-ROBOFLOW_VERSION = 10
+ROBOFLOW_VERSION = 12
 CONF_THRESHOLD = 0.35
 IOU_THRESHOLD = 0.4
 
@@ -19,9 +29,119 @@ client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 rf = Roboflow(api_key=ROBOFLOW_API_KEY)
 model = rf.workspace().project(ROBOFLOW_PROJECT).version(ROBOFLOW_VERSION).model
 
-CROP_DIR = "/data/cropped_labels"
+CROP_DIR = "cropped_labels"
 os.makedirs(CROP_DIR, exist_ok=True)
 
+# ================= CLASSIFIER =================
+
+CLASS_NAMES = ['BB', 'FF', 'P']
+
+classifier_model = models.efficientnet_b0(weights=None)
+in_features = classifier_model.classifier[1].in_features
+
+classifier_model.classifier[1] = nn.Sequential(
+    nn.Dropout(p=0.3, inplace=True),
+    nn.Linear(in_features, 3)
+)
+
+classifier_model.load_state_dict(torch.load('breaker_classifier.pth', map_location='cpu'))
+classifier_model.eval()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+classifier_model.to(device)
+
+# transform
+type_transform = A.Compose([
+    A.Resize(224, 224),
+    A.Normalize(mean=(0.485, 0.456, 0.406),
+                std=(0.229, 0.224, 0.225)),
+    ToTensorV2(),
+])
+
+def predict_breaker_type(image):
+    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    augmented = type_transform(image=img)
+    tensor = augmented['image'].unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = classifier_model(tensor)
+        probs = torch.softmax(outputs, dim=1)
+        conf, pred = torch.max(probs, 1)
+
+    return CLASS_NAMES[pred.item()], float(conf.item())
+
+
+# ================= EQUIPMENT TYPE CLASSIFIER =================
+
+
+
+# ================= DEVICE =================
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ================= CLASS NAMES (MUST MATCH TRAINING) =================
+EQUIPMENT_CLASS_NAMES = ['ACB', 'E', 'EM', 'EMDU', 'EMMDU', 'M', 'MDU']
+
+# ================= TRANSFORM (EXACT SAME AS TRAINING) =================
+val_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+# ================= LOAD MODEL =================
+equipment_model = models.efficientnet_b0(pretrained=False)
+
+# IMPORTANT: same classifier as training
+equipment_model.classifier[1] = nn.Linear(
+    equipment_model.classifier[1].in_features,
+    len(EQUIPMENT_CLASS_NAMES)
+)
+
+# Load weights
+equipment_model.load_state_dict(
+    torch.load("efficientnet_breaker.pth", map_location=device)
+)
+
+equipment_model = equipment_model.to(device)
+equipment_model.eval()
+
+# ================= PREDICTION FUNCTION =================
+def predict_equipment_type(image):
+    """
+    Input: OpenCV image (BGR)
+    Output:
+        best_class -> predicted label
+        best_conf  -> confidence
+        prob_dict  -> all class probabilities
+    """
+
+    # ---- Convert OpenCV → PIL (CRITICAL STEP) ----
+    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    img = Image.fromarray(img)
+
+    # ---- Apply SAME transform as training ----
+    img = val_transform(img).unsqueeze(0).to(device)
+
+    # ---- Inference ----
+    with torch.no_grad():
+        outputs = equipment_model(img)
+        probs = F.softmax(outputs, dim=1)
+
+    # ---- Convert probabilities ----
+    prob_dict = {
+        EQUIPMENT_CLASS_NAMES[i]: float(probs[0][i].item())
+        for i in range(len(EQUIPMENT_CLASS_NAMES))
+    }
+
+    # ---- Get best class ----
+    pred_idx = torch.argmax(probs, dim=1).item()
+    best_class = EQUIPMENT_CLASS_NAMES[pred_idx]
+    best_conf = prob_dict[best_class]
+
+    return best_class, best_conf, prob_dict
+
 # ================= CONSTANTS =================
 
 CIRCUIT_PATTERN = r"(?:\d+L\d+-\d+|S\d+)"
@@ -30,13 +150,37 @@ DEFAULT_BREAKING_CAPACITY = "85"
 
 VALID_AF_VALUES = {"50","63","100","125","160","250","400","630"}
 
+VALID_OPTIONS = [
+    "AL",
+    "AX",
+    "PAL",
+    "EAL",
+    "SHT",
+    "AL+AX",
+    "AL+AX+PAL+EAL"
+]
+
+# ================= LABEL MAPPING (ADD HERE) =================
+LABEL_MAP = {
+    "manufacture name": "Manufacture Name",
+    "load name": "Load Name",
+    "breaking capacity": "Breaking Capacity",
+    "af": "AF",
+    "at": "AT",
+    "option": "Option",
+    "circuit name": "Circuit Name"
+}
+
 SPEC_JP = {
     "Manufacture Name": "メーカー",
     "Circuit Name": "回路番号",
     "Load Name": "負荷名称",
     "Breaking Capacity": "遮断容量",
     "AF": "フレーム(AF)",
-    "AT": "トリップ(AT)"
+    "AT": "トリップ(AT)",
+    "Option": "オプション",
+    "Type": "タイプ",
+    "Equipment Type": "機器種別"
 }
 
 MANUFACTURER_JP_MAP = {
@@ -47,6 +191,7 @@ MANUFACTURER_JP_MAP = {
     "LS ELECTRIC": "LS ELECTRIC"
 }
 
+# ✅ STRICT WHITELIST
 KNOWN_MANUFACTURERS = {
     "MITSUBISHI ELECTRIC",
     "SIEMENS",
@@ -75,7 +220,7 @@ def crop_with_optional_expand(img, x1, y1, x2, y2, label):
     return img[max(0,y1):min(h,y2), max(0,x1):min(w,x2)]
 
 def upscale(img):
-    return cv2.resize(img, None, fx=2.5, fy=2.5, interpolation=cv2.INTER_CUBIC)
+    return cv2.resize(img, None, fx=1.5, fy=1.5, interpolation=cv2.INTER_CUBIC)
 
 def rotate_image(img, a):
     if a == 90: return cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
@@ -109,6 +254,73 @@ def enhance_AT(img):
     sharp = cv2.addWeighted(img, 1.5, blur, -0.5, 0)
     return sharp
 
+def enhance_option(img):
+    # upscale for tiny text
+    img = cv2.resize(img, None, fx=2.2, fy=2.2, interpolation=cv2.INTER_CUBIC)
+
+    # --- HSV boost (brightness + saturation) ---
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+
+    # brighten dark labels
+    hsv[:,:,2] = cv2.add(hsv[:,:,2], 55)
+
+    # increase saturation so faded letters appear
+    hsv[:,:,1] = cv2.add(hsv[:,:,1], 35)
+
+    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+
+    # --- gamma correction (critical for dark plastic labels) ---
+    gamma = 1.6
+    inv = 1.0 / gamma
+    table = np.array([(i/255.0)**inv * 255 for i in range(256)]).astype("uint8")
+    img = cv2.LUT(img, table)
+
+    # grayscale for OCR
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+    # denoise
+    gray = cv2.fastNlMeansDenoising(gray, h=10)
+
+    # strong CLAHE (local contrast)
+    clahe = cv2.createCLAHE(4.0, (8,8))
+    gray = clahe.apply(gray)
+
+    # strong sharpen → reveals dark letters
+    blur = cv2.GaussianBlur(gray, (0,0), 1.3)
+    sharp = cv2.addWeighted(gray, 1.9, blur, -0.9, 0)
+
+    return cv2.cvtColor(sharp, cv2.COLOR_GRAY2BGR)
+
+def normalize_option_text(text):
+    if not text:
+        return ""
+
+    t = text.upper().strip()
+
+    # ---- orientation / OCR confusion map
+    confusion_map = {
+        "7": "A",   # upside-down A
+        "V": "L",   # upside-down L
+        "1": "L",
+        "|": "L",
+        "I": "L",
+        "Y": "L",
+        "4": "A"
+    }
+
+    fixed = "".join(confusion_map.get(c, c) for c in t)
+
+    # direct match
+    if fixed in VALID_OPTIONS:
+        return fixed
+
+    # fuzzy match (core industrial step)
+    m = get_close_matches(fixed, VALID_OPTIONS, n=1, cutoff=0.4)
+    if m:
+        return m[0]
+
+    return fixed
+
 def save_crop(label, img):
     cv2.imwrite(os.path.join(CROP_DIR, f"{label}.jpg"), img)
 
@@ -142,19 +354,37 @@ def normalize_for_compare(t):
 
 def gpt_single_ocr(label, img):
     b64 = img_to_base64(img)
+
     rules = {
         "Manufacture Name": "Read manufacturer name in English only.",
         "Circuit Name": "Read the FULL text exactly as printed.",
         "Load Name": "Read exact text.",
         "AF": "Read the FULL text exactly as printed.",
         "AT": "Read the FULL text exactly as printed.",
-        "Breaking Capacity": "Read the FULL text exactly as printed."
+        "Breaking Capacity": "Read the FULL text exactly as printed.",
+
+        # ⭐ NEW — OPTION PROMPT
+        "Option": (
+            "Industrial breaker OPTION label. "
+            "Text is short (1–10 characters). "
+            "Image may be dark, faint, small, or rotated. "
+            "Return ONLY the exact printed text. "
+            "If unreadable return empty."
+        )
     }
 
     r = client.chat.completions.create(
         model="gpt-5.2",
         messages=[
-            {"role":"system","content":"You are a strict OCR engine."},
+            {
+                "role": "system",
+                "content": (
+                    "You are an industrial electrical label OCR engine. "
+                    "Extract text exactly as printed. "
+                    "Text may be tiny, dark, or rotated. "
+                    "Return only text."
+                )
+            },
             {"role":"user","content":[
                 {"type":"text","text":rules[label]},
                 {"type":"image_url","image_url":{"url":f"data:image/jpeg;base64,{b64}"}}
@@ -162,19 +392,22 @@ def gpt_single_ocr(label, img):
         ],
         temperature=0
     )
+
     return r.choices[0].message.content.strip()
 
+
 # ================= OCR CORE =================
 
 def gpt_ocr(label, img):
 
-    # ================= MANUFACTURER ================
+    # ================= MANUFACTURER =================
     if label in ["Manufacture Name","Load Name"]:
         img = enhance(img)
         save_crop(label, img)
         t = gpt_single_ocr(label, img)
         return normalize_manufacturer(t) if label=="Manufacture Name" else remove_spaces_only(t)
 
+    # ================= BREAKING CAPACITY =================
     if label == "Breaking Capacity":
         img = enhance_breaking_capacity(img)
         t = gpt_single_ocr(label, img)
@@ -184,14 +417,19 @@ def gpt_ocr(label, img):
 
 
     # ========= ROTATION BASED LABELS =========
+
     best_text = ""
     best_score = -1
     best_img = None
 
-    base = enhance(img)
-    for a in [0, 90]:
-        rimg = rotate_image(base, a)
+    # ⭐ IMPORTANT — choose preprocessing based on label
+    if label == "Option":
+        base = enhance_option(img)
+    else:
+        base = enhance(img)
 
+    for a in [0, 90, 180, 270]:
+        rimg = rotate_image(base, a)
 
         try:
             t = gpt_single_ocr(label, rimg)
@@ -248,6 +486,26 @@ def gpt_ocr(label, img):
             continue
 
 
+        # ================= OPTION ⭐ NEW =================
+        elif label == "Option":
+
+            # ⭐ IMPORTANT — normalize BEFORE scoring
+            candidate = normalize_option_text(clean)
+
+            # Option text usually short
+            if len(candidate) <= 10:
+                score += 40
+
+            # prefer valid option hits (strong signal)
+            if candidate in VALID_OPTIONS:
+                score += 120
+
+            # prefer alphabetic (real options are alphabetic)
+            if re.search(r"[A-Za-z]", candidate):
+                score += 20
+
+            score += len(candidate)
+
         # Track best candidate
         if score > best_score:
             best_score = score
@@ -262,7 +520,6 @@ def gpt_ocr(label, img):
     return ""
 
 
-
 # ================= EXCEL VERIFICATION =================
 
 def normalize_header(s):
@@ -306,23 +563,13 @@ def verify_excel(excel, det):
 
     if hdr is None:
         return pd.DataFrame([
-            ["Excel", "", "エラー", "項目が見つかりません。"]
+            ["Excel", "", "エラー", "ヘッダー行が見つかりません。"]
         ], columns=["仕様","検出値","Excelに存在？","備考"])
 
     df = raw.iloc[hdr+1:].copy()
     df.columns = raw.iloc[hdr]
     df.dropna(how="all", inplace=True)
 
-    def normalize_header(s):
-        return str(s).replace("\n","").replace(" ","")
-
-    def find_column(df, keys):
-        for c in df.columns:
-            for k in keys:
-                if k in normalize_header(c):
-                    return c
-        return None
-
     ccol = find_column(df, ["回路番号","回路"])
 
     if ccol is None:
@@ -343,10 +590,26 @@ def verify_excel(excel, det):
 
     rows = []
 
+    # ⭐ LOOP THROUGH ALL FIELDS (NOW INCLUDES TYPE)
     for k, jp in SPEC_JP.items():
 
         detected_value = det.get(k, "").strip()
-        col = find_column(df, [jp.replace("(","").replace(")",""), jp[:2]])
+
+        # ---- column search keys ----
+        keys = [jp.replace("(","").replace(")",""), jp[:2]]
+
+        # ⭐ OPTION SUPPORT
+        if k == "Option":
+            keys += ["オプション", "オプシ", "ション"]
+
+        # ⭐ TYPE SUPPORT (NEW 🔥)
+        if k == "Type":
+            keys += ["タイプ"]
+
+        if k == "Equipment Type":
+            keys += ["機器 種別"]
+
+        col = find_column(df, keys)
 
         if col is None:
             rows.append([
@@ -359,12 +622,19 @@ def verify_excel(excel, det):
 
         excel_value = str(target[col])
 
-        if k=="Manufacture Name":
-            for eng,jpn in MANUFACTURER_JP_MAP.items():
-                if jpn in excel_value:
-                    excel_value = eng
-                    break
+        if k == "Manufacture Name":
+            detected_jp = MANUFACTURER_JP_MAP.get(detected_value, detected_value)
+            excel_jp = excel_value
+            ok = normalize_for_compare(detected_jp) == normalize_for_compare(excel_jp)
+            rows.append([
+                jp,
+                detected_jp,   # 👈 show Japanese here
+                "YES" if ok else "NO",
+                "" if ok else f"Excel値: {excel_jp}"
+            ])
+            continue
 
+        # not detected in panel
         if not detected_value:
             rows.append([
                 jp,
@@ -386,11 +656,10 @@ def verify_excel(excel, det):
     return pd.DataFrame(rows,columns=["仕様","検出値","Excelに存在？","備考"])
 
 
-
-
 # ================= PIPELINE & UI =================
 
-def bbox_area(p): return p["width"] * p["height"]
+def bbox_area(p): 
+    return p["width"] * p["height"]
 
 def run_pipeline(image, excel):
 
@@ -402,7 +671,9 @@ def run_pipeline(image, excel):
         return None, pd.DataFrame(), pd.DataFrame(), None, \
             "⚠️ **Please upload the breaker panel image before running verification.**"
 
+    # ================= DETECTION =================
     img = prepare_for_roboflow(image)
+
     preds = model.predict(
         img,
         confidence=int(CONF_THRESHOLD*100),
@@ -410,37 +681,68 @@ def run_pipeline(image, excel):
     ).json()["predictions"]
 
     vis = img.copy()
-    det={}
-    best_boxes={}
+    det = {}
+    best_boxes = {}
 
+    # ================= SELECT BEST BOX =================
     for p in preds:
-        lab=p["class"]
+        raw_lab = p["class"]
+        lab = LABEL_MAP.get(raw_lab.lower(), raw_lab)
+
         if lab not in best_boxes:
-            best_boxes[lab]=p
+            best_boxes[lab] = p
         else:
-            if lab=="Circuit Name":
-                if bbox_area(p)<bbox_area(best_boxes[lab]):
-                    best_boxes[lab]=p
+            if lab == "Circuit Name":
+                if bbox_area(p) < bbox_area(best_boxes[lab]):
+                    best_boxes[lab] = p
             else:
-                if p["confidence"]>best_boxes[lab]["confidence"]:
-                    best_boxes[lab]=p
+                if p["confidence"] > best_boxes[lab]["confidence"]:
+                    best_boxes[lab] = p
+
+    # ================= PARALLEL OCR =================
+    def process_label(item):
+        lab, p = item
+
+        x, y, w, h = map(int, [p["x"], p["y"], p["width"], p["height"]])
+        x1, y1, x2, y2 = x - w//2, y - h//2, x + w//2, y + h//2
+
+        roi = upscale(crop_with_optional_expand(img, x1, y1, x2, y2, lab))
+
+        try:
+            value = gpt_ocr(lab, roi)
+        except Exception as e:
+            value = ""
+            print(f"OCR Error for {lab}: {e}")
 
-    for lab,p in best_boxes.items():
-        x,y,w,h=map(int,[p["x"],p["y"],p["width"],p["height"]])
-        x1,y1,x2,y2=x-w//2,y-h//2,x+w//2,y+h//2
-        cv2.rectangle(vis,(x1,y1),(x2,y2),(0,255,0),2)
-        roi = upscale(crop_with_optional_expand(img,x1,y1,x2,y2,lab))
-        det[lab]=gpt_ocr(lab,roi)
+        return lab, value, (x1, y1, x2, y2)
 
-    ocr_df=pd.DataFrame(det.items(),columns=["Field","Extracted Text"])
-    verify_df=verify_excel(excel,det)
+    # 🔥 Parallel execution (IMPORTANT)
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        results = list(executor.map(process_label, best_boxes.items()))
 
-    out="verification_result.xlsx"
-    with pd.ExcelWriter(out,engine="openpyxl") as w:
-        ocr_df.to_excel(w,"OCR_Output",index=False)
-        verify_df.to_excel(w,"Verification",index=False)
+    # ================= COLLECT RESULTS =================
+    for lab, value, (x1, y1, x2, y2) in results:
+        det[lab] = value
+        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
 
-    return vis,ocr_df,verify_df,out, ""
+    # ================= TYPE CLASSIFICATION =================
+    pred_type, conf = predict_breaker_type(img)
+    det["Type"] = pred_type
+
+    eq_type, eq_conf, _ = predict_equipment_type(img)
+    det["Equipment Type"] = eq_type
+
+    # ================= DATAFRAMES =================
+    ocr_df = pd.DataFrame(det.items(), columns=["Field", "Extracted Text"])
+    verify_df = verify_excel(excel, det)
+
+    # ================= SAVE =================
+    out = "verification_result.xlsx"
+    with pd.ExcelWriter(out, engine="openpyxl") as w:
+        ocr_df.to_excel(w, "OCR_Output", index=False)
+        verify_df.to_excel(w, "Verification", index=False)
+
+    return vis, ocr_df, verify_df, out, ""
 
 apple_dark_pink_css = """
 @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600&display=swap');
@@ -491,34 +793,6 @@ input, textarea, select {
 }
 """
 
-import gradio as gr
-
-apple_dark_pink_css = """
-@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600&display=swap');
-.gradio-container {
-    background: #0f1115;
-    font-family: 'Outfit', sans-serif;
-}
-h1 { color: #f9fafb; font-weight: 600; }
-h2, h3 { color: #e5e7eb; font-weight: 500; }
-.gr-box {
-    background: #161a22;
-    border-radius: 16px;
-    padding: 12px;
-}
-button.primary {
-    background: #f472b6 !important;
-    color: #020617 !important;
-    border-radius: 12px;
-    font-weight: 500;
-}
-button.primary:hover {
-    background: #ec4899 !important;
-}
-input, textarea, select {
-    border-radius: 10px !important;
-}
-"""
 with gr.Blocks(
     theme=gr.themes.Soft(primary_hue="pink"),
     css=apple_dark_pink_css
@@ -545,4 +819,4 @@ with gr.Blocks(
         [img_out, t1, t2, f, status_msg]
     )
 
-demo.launch()
+demo.launch()