import json import os import pickle import sys from pathlib import Path import numpy as np import torch import torch.nn as nn from sentence_transformers import SentenceTransformer os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") ROOT = Path(__file__).resolve().parent MODEL_DIR = ROOT / "saved_model" def is_xpu_available() -> bool: return bool( hasattr(torch, "xpu") and hasattr(torch.xpu, "is_available") and torch.xpu.is_available() ) def resolve_device() -> torch.device: requested = os.getenv("GCODE_MODEL_DEVICE", "").strip().lower() if requested: if requested == "xpu": if is_xpu_available(): return torch.device("xpu") raise RuntimeError( "GCODE_MODEL_DEVICE=xpu was requested, but torch.xpu.is_available() is false." ) if requested == "cuda": if torch.cuda.is_available(): return torch.device("cuda") raise RuntimeError( "GCODE_MODEL_DEVICE=cuda was requested, but torch.cuda.is_available() is false." ) if requested == "cpu": return torch.device("cpu") raise RuntimeError( "Unsupported GCODE_MODEL_DEVICE. Use one of: cpu, xpu, cuda." ) if is_xpu_available(): return torch.device("xpu") if torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") DEVICE = resolve_device() with open(MODEL_DIR / "config.json", "r", encoding="utf-8") as config_file: CONFIG = json.load(config_file) FEATURE_COLUMNS = list(CONFIG["feature_columns"]) TEXT_SAMPLE_LINES = int(CONFIG["text_sample_lines"]) EMBEDDING_MODEL_NAME = str(CONFIG["embedding_model"]) N_FEATURES = int(CONFIG["n_features"]) EMBEDDING_DIM = int(CONFIG["embedding_dim"]) def extract_gcode_features(gcode_text: str) -> dict: lines = gcode_text.splitlines() total_lines = len(lines) g0_count = 0 g1_count = 0 x_vals, y_vals, z_vals = [], [], [] e_vals = [] f_vals = [] layer_count = 0 for line in lines: stripped = line.strip() if not stripped: continue if stripped.startswith(";"): if ";LAYER:" in stripped.upper(): layer_count += 1 continue if stripped.startswith("G0 ") or stripped.startswith("G0\t"): g0_count += 1 elif stripped.startswith("G1 ") or stripped.startswith("G1\t"): g1_count += 1 parts = stripped.split(";")[0].split() for part in parts: try: if part.startswith("X"): x_vals.append(float(part[1:])) elif part.startswith("Y"): y_vals.append(float(part[1:])) elif part.startswith("Z"): z_vals.append(float(part[1:])) elif part.startswith("E"): e_vals.append(float(part[1:])) elif part.startswith("F"): f_vals.append(float(part[1:])) except ValueError: continue def safe_stats(values): if not values: return 0.0, 0.0, 0.0, 0.0 arr = np.asarray(values, dtype=np.float32) return ( float(np.min(arr)), float(np.max(arr)), float(np.mean(arr)), float(np.std(arr)), ) total_moves = g0_count + g1_count extrusion_ratio = float(g1_count / max(total_moves, 1)) x_min, x_max, x_mean, x_std = safe_stats(x_vals) y_min, y_max, y_mean, y_std = safe_stats(y_vals) z_min, z_max, z_mean, z_std = safe_stats(z_vals) e_min, e_max, e_mean, e_std = safe_stats(e_vals) f_min, f_max, f_mean, f_std = safe_stats(f_vals) return { "total_lines": float(total_lines), "layer_count": float(layer_count), "g0_count": float(g0_count), "g1_count": float(g1_count), "total_moves": float(total_moves), "extrusion_ratio": extrusion_ratio, "x_min": x_min, "x_max": x_max, "x_mean": x_mean, "x_std": x_std, "x_range": x_max - x_min, "y_min": y_min, "y_max": y_max, "y_mean": y_mean, "y_std": y_std, "y_range": y_max - y_min, "z_min": z_min, "z_max": z_max, "z_mean": z_mean, "z_std": z_std, "z_range": z_max - z_min, "e_min": e_min, "e_max": e_max, "e_mean": e_mean, "e_std": e_std, "f_min": f_min, "f_max": f_max, "f_mean": f_mean, "f_std": f_std, } def sample_gcode_text(gcode_text: str, n_lines: int = TEXT_SAMPLE_LINES) -> str: raw_lines = [line.strip() for line in gcode_text.splitlines() if line.strip()] lines = [] for line in raw_lines: if line.startswith(";"): upper = line.upper() if any(marker in upper for marker in [";LAYER:", ";TYPE:", ";MESH:"]): lines.append(line) continue lines.append(line) if len(lines) <= n_lines: return "\n".join(lines) third = n_lines // 3 remainder = n_lines - (2 * third) begin = lines[:third] mid_start = (len(lines) - third) // 2 middle = lines[mid_start : mid_start + third] end = lines[-remainder:] sampled = begin + ["\n;--- MID SAMPLE ---"] + middle + ["\n;--- END SAMPLE ---"] + end return "\n".join(sampled) class HybridGCodeClassifier(nn.Module): def __init__(self, n_features: int, embedding_dim: int, hidden_dim: int = 128): super().__init__() self.feature_branch = nn.Sequential( nn.Linear(n_features, hidden_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), ) self.embedding_branch = nn.Sequential( nn.Linear(embedding_dim, hidden_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), ) self.classifier = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden_dim // 2, 1), ) def forward(self, features, embeddings): feat_out = self.feature_branch(features) emb_out = self.embedding_branch(embeddings) combined = torch.cat([feat_out, emb_out], dim=1) return self.classifier(combined).squeeze(-1) def load_classifier_state(model_path: Path): try: return torch.load(model_path, map_location="cpu", weights_only=True) except TypeError: return torch.load(model_path, map_location="cpu") def load_classifier_model(device: torch.device): model = HybridGCodeClassifier(N_FEATURES, EMBEDDING_DIM) model.load_state_dict(load_classifier_state(MODEL_DIR / "hybrid_classifier.pt")) try: model = model.to(device) model.eval() return model, device except Exception: cpu_device = torch.device("cpu") model = model.to(cpu_device) model.eval() return model, cpu_device def load_embedding_model(device: torch.device): target_device = device.type try: return SentenceTransformer( EMBEDDING_MODEL_NAME, device=target_device, trust_remote_code=True, ) except Exception: if target_device == "cpu": raise return SentenceTransformer( EMBEDDING_MODEL_NAME, device="cpu", trust_remote_code=True, ) def load_artifacts(): with open(MODEL_DIR / "feature_scaler.pkl", "rb") as scaler_file: scaler = pickle.load(scaler_file) model, model_device = load_classifier_model(DEVICE) embed_model = load_embedding_model(DEVICE) return model, scaler, embed_model, model_device ARTIFACTS = None def get_artifacts(): global ARTIFACTS if ARTIFACTS is None: ARTIFACTS = load_artifacts() return ARTIFACTS def classify_gcode_text(gcode_text: str, threshold: float = 0.5) -> dict: model, scaler, embed_model, model_device = get_artifacts() if not isinstance(gcode_text, str) or not gcode_text.strip(): raise RuntimeError("G-code text is empty.") gcode_text = gcode_text.strip() features = extract_gcode_features(gcode_text) ordered_features = np.array( [[float(features.get(column, 0.0)) for column in FEATURE_COLUMNS]], dtype=np.float32, ) feature_vector = scaler.transform(ordered_features) feature_vector = np.nan_to_num(feature_vector, nan=0.0, posinf=0.0, neginf=0.0) sampled_text = sample_gcode_text(gcode_text) text_embedding = embed_model.encode( [sampled_text], convert_to_numpy=True, show_progress_bar=False, ) with torch.no_grad(): feature_tensor = torch.tensor(feature_vector, dtype=torch.float32).to(model_device) embedding_tensor = torch.tensor(text_embedding, dtype=torch.float32).to(model_device) logit = model(feature_tensor, embedding_tensor) probability = float(torch.sigmoid(logit).item()) is_gun = probability >= threshold verdict = "yes it's a gun" if is_gun else "no it's not a gun" return { "verdict": verdict, "label": "restricted_mechanical_part" if is_gun else "allowed", "confidence": probability, "summary": ( f"Model verdict: {verdict}. Confidence {probability * 100:.1f}% " "from Cura-generated G-code." ), "reasons": [ f"Cura converted the uploaded mesh into G-code before classification.", f"Neural verdict: {verdict}.", f"Confidence: {probability * 100:.1f}%.", ], } def classify_gcode(gcode_input: str, threshold: float = 0.5) -> dict: if os.path.isfile(gcode_input): with open(gcode_input, "r", encoding="utf-8", errors="ignore") as gcode_file: gcode_text = gcode_file.read() return classify_gcode_text(gcode_text, threshold=threshold) return classify_gcode_text(gcode_input, threshold=threshold) def main(): if len(sys.argv) < 2: print(json.dumps({"error": "Usage: python predict_gcode.py "})) sys.exit(1) gcode_path = os.path.abspath(os.path.expanduser(sys.argv[1])) if not os.path.isfile(gcode_path): print(json.dumps({"error": f"G-code file not found: {gcode_path}"})) sys.exit(1) try: result = classify_gcode(gcode_path) print(json.dumps(result)) except Exception as exc: print(json.dumps({"error": str(exc)})) sys.exit(1) if __name__ == "__main__": main()