app.py

#!/usr/bin/env python3
"""
CircuitGlyph — Flask Backend
"""

import os, base64, tempfile, logging
import torch
import timm
import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from flask import Flask, request, jsonify, send_from_directory
from werkzeug.utils import secure_filename
from torchinfo import summary
from huggingface_hub import hf_hub_download

# ─── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ─── CONFIG — edit these paths ─────────────────────────────────────────────────
TRAIN_FOLDER_PATH = "./master_classification"
FRONT_BINARY_PATH = hf_hub_download(
    repo_id="ssayed122/circuitglyph-models",
    filename="front_binary.pth"
)
CHECKPOINT_PATH = hf_hub_download(
    repo_id="ssayed122/circuitglyph-models",
    filename="masterCNN.pth"
)
CHECKPOINT_PATH_VIT = hf_hub_download(
    repo_id="ssayed122/circuitglyph-models",
    filename="masterViT.pth"
)
SUB_MODEL_FOLDER     = "./models"                                      # folder with e.g. 555-timer.pth
SUB_TRAIN_FOLDER_PATH = "./sub_classification"

MODEL_NAME_EFFNET    = 'efficientnet_b1'
MODEL_NAME_FRONT_BIN = 'efficientnet_b2'
MODEL_NAME_VIT       = 'vit_small_patch16_224'
SUB_MODEL_NAME_EFFNET = 'efficientnet_b2'
NUM_CLASSES          = 51
INPUT_SIZE           = 384
INPUT_SIZE_VIT       = 224
INPUT_SIZE_FRONT     = 380

DEVICE = torch.device("mps" if torch.backends.mps.is_available() else
                       "cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {DEVICE}")

# ─── Sub-model config ──────────────────────────────────────────────────────────
FRONT_BIN_CLASSES = ['circuit','not-a-circuit']
SUB_CLASSES = [
    '555-timer','am-fm-sigma-delta-modulator','antilog-log-amplifier',
    'bjt-amplifier','cascode-darlington-amplifier','differentiator-integrator',
    'digital-adder','flip-flop-latch','frequency-divider-multiplier','memory-cell',
    'mos-amplifier','multiplexer-demultiplexer','operational-amplifier','oscillator',
    'power-amplifier','power-converter','protection-safety-circuit',
    'specialized-amplifier','active-filter','passive-filter'
]
SUB_CLASSES_MODEL_FILE = [
    '555-timer','am-fm-sigma-delta-modulator','antilog-log-amplifier',
    'bjt-amplifier','cascode-darlington-amplifier','differentiator-integrator',
    'digital-adder','flip-flop-latch','frequency-divider-multiplier','memory-cell',
    'mos-amplifier','multiplexer-demultiplexer','operational-amplifier','oscillator',
    'power-amplifier','power-converter','protection-safety-circuit',
    'specialized-amplifier','filter','filter'
]
SUB_CLASSES_DESC = [
    'a 555 timer','AM/FM or sigma-delta modulator','a log or an antilog amplifier',
    'a bjt amplifier','a cascode or a darlington amplifier','a differentiator or an integrator',
    'a digital adder','a flip-flop/latch','a frequency divider or multiplier','a memory cell',
    'a mos amplifier','a multiplexer or demultiplexer','an operational amplifier','an oscillator',
    'a power amplifier','a power converter','a protection safety',
    'a specialized amplifier','a filter','a filter'
]
SUB_NUM_CLASSES = [2,3,2,3,3,3,3,6,2,2,2,2,6,11,4,11,4,11,9,9]

# ─── Transforms ───────────────────────────────────────────────────────────────
inference_transform = A.Compose([
    A.Resize(INPUT_SIZE, INPUT_SIZE),
    A.ToFloat(max_value=255.0),
    ToTensorV2(),
])
inference_transform_vit = A.Compose([
    A.Resize(INPUT_SIZE_VIT, INPUT_SIZE_VIT),
    A.ToFloat(max_value=255.0),
    ToTensorV2(),
])
inference_transform_front = A.Compose([
    A.Resize(INPUT_SIZE_FRONT, INPUT_SIZE_FRONT),
    A.ToFloat(max_value=255.0),
    ToTensorV2(),
])

# ─── Helpers ───────────────────────────────────────────────────────────────────
def load_class_names(folder):
    if not os.path.exists(folder):
        logger.warning(f"Folder not found: {folder}")
        return None
    names = sorted([d for d in os.listdir(folder) if os.path.isdir(os.path.join(folder, d))])
    logger.info(f"Loaded {len(names)} class names from {folder}")
    return names

def preprocess_image(src_path, dest_path, target=640, pad_value=255):
    img = cv2.imread(src_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise ValueError(f"Could not read image: {src_path}")
    img = cv2.GaussianBlur(img, (3, 3), 0)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    img = clahe.apply(img)
    img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                 cv2.THRESH_BINARY_INV, 31, 7)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
    img = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
    h, w = img.shape
    scale = target / max(h, w)
    nh, nw = int(h*scale), int(w*scale)
    img = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA)
    dw, dh = target-nw, target-nh
    img = cv2.copyMakeBorder(img, dh//2, dh-dh//2, dw//2, dw-dw//2,
                              cv2.BORDER_CONSTANT, value=pad_value)
    cv2.imwrite(dest_path, img)

def load_model(model_name, checkpoint_path, num_classes):
    model = timm.create_model(model_name, pretrained=False,
                               num_classes=num_classes,
                               drop_rate=0.30, drop_path_rate=0.25)
    model = model.to(DEVICE)
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
    ckpt = torch.load(checkpoint_path, map_location=DEVICE)
    if isinstance(ckpt, dict):
        sd = ckpt.get('model_state_dict') or ckpt.get('state_dict') or ckpt
        model.load_state_dict(sd)
    else:
        model = ckpt.to(DEVICE)
    model.eval()
    logger.info(f"Loaded {model_name} from {checkpoint_path}")
    logger.info(sum(p.numel() for p in model.parameters()))
    return model

def load_sub_model(model_file_name):
    path = hf_hub_download(
        repo_id="ssayed122/circuitglyph-models",
        filename=f"{model_file_name}.pth"
    )
    ckpt = torch.load(path, map_location=DEVICE)
    sd = ckpt.get('model_state_dict') or ckpt.get('state_dict') or ckpt if isinstance(ckpt, dict) else ckpt
    num_sub = sd[list(sd.keys())[-2]].shape[0]
    sub_model = timm.create_model(SUB_MODEL_NAME_EFFNET, pretrained=False,
                                   num_classes=num_sub, drop_rate=0.30, drop_path_rate=0.25)
    sub_model.load_state_dict(sd)
    sub_model = sub_model.to(DEVICE)
    sub_model.eval()
    sub_folder = os.path.join(SUB_TRAIN_FOLDER_PATH, f"{model_file_name}")
    sub_class_names = load_class_names(sub_folder)
    return sub_model, sub_class_names

def front_binary(image_path):
    model = load_model(MODEL_NAME_FRONT_BIN, FRONT_BINARY_PATH, 2)
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    tensor = inference_transform_front(image=img)['image'].unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        probs = torch.softmax(model(tensor), dim=1)[0]
    conf, idx = torch.max(probs, 0)
    label = FRONT_BIN_CLASSES[idx.item()] if FRONT_BIN_CLASSES else f"Class {idx.item()}"
    return label, round(conf.item(), 4)
    
def predict_image(image_path, model, transform, class_names):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    tensor = transform(image=img)['image'].unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        probs = torch.softmax(model(tensor), dim=1)[0]
    conf, idx = torch.max(probs, 0)
    label = class_names[idx.item()] if class_names else f"Class {idx.item()}"
    k = min(5, len(class_names) if class_names else NUM_CLASSES)
    top5_prob, top5_idx = torch.topk(probs, k)
    top5 = [{"name": class_names[i.item()] if class_names else f"Class {i.item()}",
              "conf": round(p.item(), 4)}
             for i, p in zip(top5_idx, top5_prob)]
    return label, round(conf.item(), 4), top5

# ─── Model cache (loaded once at startup) ─────────────────────────────────────
_models = {}

def get_models():
    if 'effnet' not in _models:
        logger.info("Loading EfficientNet model…")
        _models['effnet'] = load_model(MODEL_NAME_EFFNET, CHECKPOINT_PATH, NUM_CLASSES)
    if 'vit' not in _models:
        logger.info("Loading ViT model…")
        _models['vit'] = load_model(MODEL_NAME_VIT, CHECKPOINT_PATH_VIT, NUM_CLASSES)
    if 'class_names' not in _models:
        _models['class_names'] = load_class_names(TRAIN_FOLDER_PATH)
    return _models['effnet'], _models['vit'], _models['class_names']

# ─── Flask app ─────────────────────────────────────────────────────────────────
app = Flask(__name__, static_folder='.')

@app.route('/')
def index():
    return send_from_directory('.', 'index.html')

@app.route("/logo.png")
def logo():
    return send_from_directory(".", "logo.png")

@app.route('/analyze', methods=['POST'])
def analyze():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400

    file = request.files['image']
    if file.filename == '':
        return jsonify({'error': 'Empty filename'}), 400
    
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, 'input.png')
        proc_path  = os.path.join(tmpdir, 'processed.png')
        file.save(input_path)

        try:
            front_pred, front_conf = front_binary(input_path)
            if front_pred == "circuit":
                # Load models
                effnet_model, vit_model, class_names = get_models()
    
                # Preprocess
                preprocess_image(input_path, proc_path)
    
                # Read processed image as base64 to send back to frontend
                with open(proc_path, 'rb') as f:
                    proc_b64 = base64.b64encode(f.read()).decode('utf-8')
    
                # Predict
                effnet_pred, effnet_conf, effnet_top5 = predict_image(proc_path, effnet_model, inference_transform, class_names)
                vit_pred,    vit_conf,    _            = predict_image(proc_path, vit_model, inference_transform_vit, class_names)
                
                # Sub-model
                sub_triggered = False
                sub_pred, sub_conf = None, None
                
                # Sub-model
                sub_triggered2 = False
                sub_pred2, sub_conf2 = None, None
                
                if effnet_pred in SUB_CLASSES:
                    sub_index = SUB_CLASSES.index(effnet_pred)
                    model_file = SUB_CLASSES_MODEL_FILE[sub_index]
                    class_desc=SUB_CLASSES_DESC[sub_index]
                    sub_model, sub_class_names = load_sub_model(model_file)
                    if sub_model is not None:
                        sub_pred, sub_conf, _ = predict_image(proc_path, sub_model, inference_transform, sub_class_names)
                        sub_triggered = True
                        sub_pred=sub_pred.replace("-"," ")
                else:
                    class_desc=str(effnet_pred).replace("-", " ")
                
                if vit_pred in SUB_CLASSES:
                    sub_index2 = SUB_CLASSES.index(vit_pred)
                    model_file2 = SUB_CLASSES_MODEL_FILE[sub_index2]
                    class_desc2=SUB_CLASSES_DESC[sub_index2]
                    sub_model2, sub_class_names2 = load_sub_model(model_file2)
                    if sub_model2 is not None:
                        sub_pred2, sub_conf2, _ = predict_image(proc_path, sub_model2, inference_transform, sub_class_names2)
                        sub_triggered2 = True
                        sub_pred2=sub_pred2.replace("-", " ")
                else:
                    class_desc2=str(vit_pred).replace("-", " ")
                        
                # Build pred_string (mirrors Python script exactly)
                pred_string = ""
                if effnet_pred==vit_pred:
                    pred_string += f"\nBoth Convolutional Neural Network (CNN)-based and Vision Transformer (ViT)-based models indicate, with {max(vit_conf, effnet_conf)*100:.2f}% confidence, that the input image depicts {class_desc2} circuit.\n"
                    if sub_triggered and sub_pred:
                        pred_string += f"\nThere's a {sub_conf*100:.2f}% chance that this circuit is a {sub_pred}.\n"
                else:
                    pred_string += f"\nA Vision Transformer (ViT) model infers with {vit_conf*100:.2f}% confidence that the input image depicts {class_desc2} circuit.\n"
                    if sub_triggered2 and sub_pred2:
                        pred_string += f"\nThere's a {sub_conf2*100:.2f}% chance that this circuit is a {sub_pred2}.\n"
                    pred_string += f"\nA Convolutional Neural Network (CNN)-based model classifies the image as {class_desc} circuit with {effnet_conf*100:.2f}% confidence.\n"
                    if sub_triggered and sub_pred:
                        pred_string += f"\nThere's a {sub_conf*100:.2f}% chance that this circuit is a {sub_pred}.\n"
                pred_string += f"\nOther possibilities include:\n"
                for i, item in enumerate(effnet_top5[1:], 1):
                    pred_string += f"{i}. {item['name'].replace('-', ' ').capitalize()} (Confidence: {item['conf']*100:.2f}%)\n"
                    
            else:
                pred_string = "This picture does not represent a circuit diagram. Please try again with a valid picture.\n"
                vit_pred = None
                vit_conf=None
                effnet_pred=None
                effnet_conf=None
                sub_triggered=None
                sub_pred=None
                sub_conf=None
                effnet_top5=None
                proc_b64=None
                
                
            return jsonify({
                'pred_string':   pred_string,
                'vit_pred':      vit_pred,
                'vit_conf':      vit_conf,
                'effnet_pred':   effnet_pred,
                'effnet_conf':   effnet_conf,
                'sub_triggered': sub_triggered,
                'sub_pred':      sub_pred,
                'sub_conf':      sub_conf,
                'top5':          effnet_top5,
                'processed_img': proc_b64,
            })

        except FileNotFoundError as e:
            logger.error(str(e))
            return jsonify({'error': str(e)}), 500
        except Exception as e:
            logger.error(f"Analysis failed: {e}", exc_info=True)
            return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    logger.info("Starting CircuitMind server on http://localhost:8080")
    app.run(host='0.0.0.0', port=7860, debug=False)