| 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 <file.gcode>"})) |
| 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() |
|
|
