#!/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)