""" DefectFill & DefectDiffu Unified Web UI - Flask Backend ========================================================= Supports both DefectFill (few-shot LoRA training) and DefectDiffu (text-guided generation). Run from project root: python defectfill_webui/app.py Requirements: pip install flask flask-cors pillow numpy """ import os import sys import json import time import shutil import threading import subprocess import traceback from pathlib import Path from datetime import datetime from PIL import Image, ImageDraw import numpy as np import uuid import io import re import importlib.util # ============================================================================= # Configuration # ============================================================================= BASE_DIR = Path(__file__).parent PROJECT_ROOT = BASE_DIR.parent # ArtiAgent_with_WebUI/ # DefectFill paths DEFECTFILL_ROOT = PROJECT_ROOT / "ArtiAgent - DefectFill" TRAIN_PY = DEFECTFILL_ROOT / 'engine' / 'DefectFill' / 'train.py' TRAINING_DATA_DIR = DEFECTFILL_ROOT / 'src' / 'data' CHECKPOINT_DIR = DEFECTFILL_ROOT / 'engine' / 'DefectFill' / 'checkpoints' TRAINED_GALLERY_DIR = DEFECTFILL_ROOT / "storage" / "trained_products_gallery" TRAINED_GALLERY_DIR.mkdir(parents=True, exist_ok=True) TEMP_IMAGE_DIR = DEFECTFILL_ROOT / 'src' / 'data' / 'tempImage' TEMP_IMAGE_DIR.mkdir(parents=True, exist_ok=True) # DefectDiffu paths DEFECTDIFFU_ROOT = PROJECT_ROOT / "ArtiAgent - DefectDiffu" DEFECTDIFFU_ENGINE = PROJECT_ROOT / "engine" / "DefectDiffu" # Web UI dirs UPLOAD_FOLDER = BASE_DIR / "uploads" OUTPUT_DIR = BASE_DIR / "outputs" CONFIG_DIR = BASE_DIR / "config" # OpenCV availability check try: import cv2 CV2_AVAILABLE = True except ImportError: CV2_AVAILABLE = False print("[WARNING] opencv-python not installed. Image alignment unavailable.") # Ensure directories exist for d in [UPLOAD_FOLDER, OUTPUT_DIR, CONFIG_DIR]: d.mkdir(parents=True, exist_ok=True) # Domain instructions config DOMAIN_CONFIG_PATH = CONFIG_DIR / "domain_instructions.json" if not DOMAIN_CONFIG_PATH.exists(): with open(DOMAIN_CONFIG_PATH, "w") as f: json.dump({"examples": {}, "active": {}}, f, indent=2) # DefectDiffu config DEFECTDIFFU_CONFIG_PATH = CONFIG_DIR / "defectdiffu_config.json" if not DEFECTDIFFU_CONFIG_PATH.exists(): with open(DEFECTDIFFU_CONFIG_PATH, "w") as f: json.dump({"ckpt_path": "", "vae_path": "", "vlm_model": "gemma3:12b"}, f, indent=2) # Setup sys.path for imports # sys.path.insert(0, str(DEFECTFILL_ROOT / 'src' / 'segment_anything')) # sys.path.insert(0, str(DEFECTFILL_ROOT / 'src')) # sys.path.insert(0, str(DEFECTFILL_ROOT / 'pipeline')) # sys.path.insert(0, str(DEFECTFILL_ROOT)) # sys.path.insert(0, str(DEFECTDIFFU_ROOT)) # sys.path.insert(0, str(DEFECTDIFFU_ROOT / 'src')) # sys.path.insert(0, str(DEFECTDIFFU_ROOT / 'pipeline')) # ============================================================================= # Dynamic Import Helpers (avoid module name collisions) # ============================================================================= def load_orchestrator_from_path(module_path, class_name="ArtiAgentOrchestrator", extra_paths=None): """ Dynamically import a class from a Python file, temporarily adding extra_paths to sys.path so that the module's internal imports resolve to the correct project. """ if not module_path.exists(): print(f"[ERROR] Module file not found: {module_path}") return None original_sys_path = sys.path[:] # save current state try: # Insert project‑specific paths at the front if extra_paths: for p in reversed(extra_paths): # reverse to keep the given order p_str = str(p) if p_str not in sys.path: sys.path.insert(0, p_str) # Also ensure the module's own directory is available module_dir = module_path.parent if module_dir not in sys.path: sys.path.insert(0, str(module_dir)) # Use a unique module name to avoid caching collisions module_id = f"_dynamic_{class_name}_{uuid.uuid4().hex[:8]}" spec = importlib.util.spec_from_file_location(module_id, str(module_path)) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) cls = getattr(module, class_name, None) if cls is None: print(f"[ERROR] Class '{class_name}' not found in {module_path}") return cls except Exception as e: print(f"[ERROR] Failed to load {module_path}: {e}") traceback.print_exc() return None finally: # Restore original sys.path to avoid leaking sys.path[:] = original_sys_path # DefectFill imports — load directly from file to avoid sys.path collision with DefectFill extra_paths_df = [ DEFECTFILL_ROOT / 'src' / 'segment_anything', DEFECTFILL_ROOT / 'src', DEFECTFILL_ROOT / 'pipeline', DEFECTFILL_ROOT, ] df_orchestrator_path = DEFECTFILL_ROOT / "src" / "artiagent_orchestrator.py" DefectFillOrchestrator = load_orchestrator_from_path( df_orchestrator_path, "ArtiAgentOrchestrator", extra_paths_df ) print("[INFO] DefectFill orchestrator is ready.", flush=True) # DefectDiffu imports — load directly from file to avoid name collision with DefectDiffu extra_paths_dd = [ DEFECTDIFFU_ROOT / 'src', DEFECTDIFFU_ROOT / "engine" / "DefectDiffu", DEFECTDIFFU_ROOT, ] dd_orchestrator_path = DEFECTDIFFU_ROOT / 'src' / "artiagent_orchestrator.py" DefectDiffuOrchestrator = load_orchestrator_from_path( dd_orchestrator_path, "ArtiAgentOrchestrator", extra_paths_dd ) print("[INFO] DefectDiffu orchestrator is ready.") def load_helper_from_path(module_path, function_name): # Use a generic loader (can reuse load_orchestrator_from_path with no class) # But for functions, you can load the module and extract the function. original_sys_path = sys.path[:] try: # Add DefectFill paths (or pass as extra_paths) sys.path.insert(0, str(DEFECTFILL_ROOT / 'src')) sys.path.insert(0, str(DEFECTFILL_ROOT)) spec = importlib.util.spec_from_file_location(f"_helper_{uuid.uuid4().hex[:8]}", str(module_path)) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return getattr(module, function_name, None) finally: sys.path[:] = original_sys_path # Usage: align_image_to_reference = load_helper_from_path( DEFECTFILL_ROOT / 'src' / 'align_image_to_reference.py', 'align_image_to_reference' ) check_product_training_status = load_helper_from_path( DEFECTFILL_ROOT / 'src' / 'detect_similar_product.py', 'check_product_training_status' ) from flask import Flask, render_template, request, jsonify, send_from_directory from flask_cors import CORS app = Flask(__name__, template_folder=str(BASE_DIR / "templates"), static_folder=str(BASE_DIR / "static")) app.config["MAX_CONTENT_LENGTH"] = 100 * 1024 * 1024 # 100MB upload limit CORS(app) # ============================================================================= # Global State # ============================================================================= training_state = { "running": False, "process": None, "object_class": None, "defect_type": None, "output_dir": None, "log_lines": [], "start_time": None, "completed": False, "error": None } training_lock = threading.Lock() # --- Generation Job Tracking --- generation_jobs = {} jobs_lock = threading.Lock() class ProgressCapture: """Captures BOTH stdout and stderr, parses [Agent] Step X: markers.""" def __init__(self, job_id): self.job_id = job_id self._buffer = "" self._step_re = re.compile(r'\[Agent\] Step (\d+):\s*(.+)') self._defect_re = re.compile(r'\[Agent\] Defect (\d+)/(\d+):') def write(self, s): self._buffer += s while '\n' in self._buffer: line, self._buffer = self._buffer.split('\n', 1) self._process_line(line) def _process_line(self, line): m = self._step_re.search(line) if m: with jobs_lock: if self.job_id in generation_jobs: generation_jobs[self.job_id]['step'] = int(m.group(1)) generation_jobs[self.job_id]['step_text'] = m.group(2).strip() d = self._defect_re.search(line) if d: with jobs_lock: if self.job_id in generation_jobs: job = generation_jobs[self.job_id] job['defect_current'] = int(d.group(1)) job['defect_total'] = int(d.group(2)) # --- ADD THIS BLOCK BELOW --- if 'unit_total' in job and job['unit_total'] > 0: defects_per_image = job['defect_total'] image_current = job['image_current'] # Formula: (Current Image - 1) * Defects per Image + Current Defect unit_current = (image_current - 1) * defects_per_image + job['defect_current'] job['unit_current'] = min(unit_current, job['unit_total']) if '[Agent] ERROR' in line or 'Traceback (most recent call last):' in line: with jobs_lock: if self.job_id in generation_jobs: generation_jobs[self.job_id]['has_error_trace'] = True def flush(self): pass def isatty(self): return False class TeeCapture: """Duplicates output to capture AND original stream.""" def __init__(self, capture, original): self.capture = capture self.original = original def write(self, s): self.capture.write(s) self.original.write(s) def flush(self): self.capture.flush() self.original.flush() def isatty(self): return False def run_generation_job(job_id, params, mode='defectfill'): """Background thread: runs generation + captures progress.""" with jobs_lock: generation_jobs[job_id]['status'] = 'running' generation_jobs[job_id]['step'] = 1 generation_jobs[job_id]['step_text'] = 'Planning defects from product description...' generation_jobs[job_id]['started_at'] = time.time() generation_jobs[job_id]['unit_total'] = len(params['clean_paths']) * params['num_defects'] generation_jobs[job_id]['unit_current'] = 0 capture = ProgressCapture(job_id) old_stdout = sys.stdout old_stderr = sys.stderr sys.stdout = TeeCapture(capture, old_stdout) sys.stderr = TeeCapture(capture, old_stderr) try: if mode == 'defectfill': orchestrator = DefectFillOrchestrator( device=params['device'], output_dir=str(params['gen_output_dir']), checkpoint_dir=str(params['ckpt_dir']), object_class=params['object_class'], defect_type=params['defect_type'] or "", valid_object_classes=params['valid_object_classes'], valid_defect_types=params['valid_defect_types'], image_size=params['image_size'], num_steps=params['num_steps'], guidance_scale=params['guidance_scale'], domain_hint=params.get('domain_hint', '') ) else: # defectdiffu orchestrator = DefectDiffuOrchestrator( device=params['device'], output_dir=str(params['gen_output_dir']), vlm_model=params.get('vlm_model', 'gemma3:12b'), defectdiffu_ckpt=params['defectdiffu_ckpt'], vae_path=params['vae_path'], image_size=params['image_size'], num_steps=params['num_steps'] ) all_output_images = [] all_grouped_results = [] total_images = len(params['clean_paths']) for img_idx, clean_path in enumerate(params['clean_paths']): with jobs_lock: generation_jobs[job_id]['image_current'] = img_idx + 1 generation_jobs[job_id]['image_total'] = total_images generation_jobs[job_id]['step_text'] = f'Image {img_idx+1}/{total_images}: Planning defects...' if mode == 'defectfill': result = orchestrator.run( product_description=params['product_desc'], image_path=clean_path, num_defects=params['num_defects'], defect_type=params['defect_type'], object_class=params['object_class'] ) else: # defectdiffu result = orchestrator.run( product_description=params['product_desc'], image_path=clean_path, num_defects=params['num_defects'], defect_type=params['defect_type'] ) source_name = Path(clean_path).name clean_rel = Path(clean_path).relative_to(BASE_DIR) clean_url = f"/api/output/image/{clean_rel}" defects_for_image = [] for r in result.get("results", []): if r.get("success") and "output_dir" in r: out_d = Path(r["output_dir"]) if out_d.exists(): mask_file = None output_file = None blended_file = None raw_file = None for img_file in sorted(out_d.glob("*.png")): fname = img_file.name.lower() if "mask" in fname: mask_file = img_file elif "blended" in fname: blended_file = img_file elif "raw_defectdiffu" in fname or "output" in fname: raw_file = img_file elif "defect" in fname and blended_file is None: output_file = img_file rel_path = img_file.relative_to(BASE_DIR) all_output_images.append({ "filename": img_file.name, "url": f"/api/output/image/{rel_path}", "defect_type": r.get("defect_type", "unknown"), "object_class": r.get("object_class", params.get('object_class', 'unknown')), "source_index": img_idx, "source_name": source_name }) # DefectDiffu has different output files if mode == 'defectdiffu': defects_for_image.append({ "defect_type": r.get("defect_type", "unknown"), "defect_label": f"Defect {len(defects_for_image) + 1}: {r.get('defect_type', 'unknown')}", "clean_image": { "filename": source_name, "url": clean_url }, "output_image": { "filename": blended_file.name if blended_file else (output_file.name if output_file else "output.png"), "url": f"/api/output/image/{(blended_file or output_file or raw_file).relative_to(BASE_DIR)}" } if (blended_file or output_file or raw_file) else None, "raw_patch_image": { "filename": raw_file.name if raw_file else "raw.png", "url": f"/api/output/image/{raw_file.relative_to(BASE_DIR)}" } if raw_file else None, "mask_image": { "filename": mask_file.name if mask_file else "mask.png", "url": f"/api/output/image/{mask_file.relative_to(BASE_DIR)}" } if mask_file else None }) elif mask_file and (output_file or blended_file): defects_for_image.append({ "defect_type": r.get("defect_type", "unknown"), "defect_label": f"Defect {len(defects_for_image) + 1}: {r.get('defect_type', 'unknown')}", "clean_image": { "filename": source_name, "url": clean_url }, "mask_image": { "filename": mask_file.name, "url": f"/api/output/image/{mask_file.relative_to(BASE_DIR)}" }, "output_image": { "filename": output_file.name if output_file else blended_file.name, "url": f"/api/output/image/{(output_file or blended_file).relative_to(BASE_DIR)}" } }) if defects_for_image: all_grouped_results.append({ "source_index": img_idx, "source_name": source_name, "defect_count": len(defects_for_image), "defects": defects_for_image }) orchestrator.cleanup() sys.stdout = old_stdout sys.stderr = old_stderr elapsed = time.time() - generation_jobs[job_id]['started_at'] with jobs_lock: generation_jobs[job_id]['status'] = 'completed' generation_jobs[job_id]['step'] = 5 generation_jobs[job_id]['step_text'] = 'Generation complete' generation_jobs[job_id]['elapsed_time'] = elapsed generation_jobs[job_id]['result'] = { 'success': True, 'grouped_results': all_grouped_results, 'output_images': all_output_images, 'experiment_id': params.get('object_class', 'defectdiffu'), 'product_type': params.get('object_class', 'defectdiffu'), 'elapsed_time': elapsed, 'results_summary': [{"type": img['defect_type'], "success": True} for img in all_output_images], 'output_dir': str(params['gen_output_dir'].relative_to(BASE_DIR)) } except Exception as e: traceback.print_exc() with jobs_lock: generation_jobs[job_id]['status'] = 'error' generation_jobs[job_id]['error'] = str(e) generation_jobs[job_id]['step_text'] = f'Error: {str(e)}' finally: sys.stdout = old_stdout sys.stderr = old_stderr # ============================================================================= # Helpers # ============================================================================= def allowed_file(filename, exts={".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}): return Path(filename).suffix.lower() in exts def get_checkpoints(): """Discover trained checkpoints: object_class -> [defect_types].""" mapping = {} if not CHECKPOINT_DIR.exists(): return mapping for obj_dir in CHECKPOINT_DIR.iterdir(): if not obj_dir.is_dir(): continue if obj_dir.name == 'checkpoints_old': continue defects = [] for defect_dir in obj_dir.iterdir(): if not defect_dir.is_dir(): continue ckpt = defect_dir / "checkpoints" / "checkpoint_final.pt" ckpt_alt = defect_dir / "checkpoint_final.pt" if ckpt.exists() or ckpt_alt.exists(): defects.append(defect_dir.name) if defects: mapping[obj_dir.name] = defects return mapping def generate_mask_from_rect(image_path, rect, output_mask_path): """Generate a binary mask from rectangle coordinates.""" img = Image.open(image_path).convert("RGB") w, h = img.size mask = Image.new("L", (w, h), 0) draw = ImageDraw.Draw(mask) x1 = max(0, int(rect["x"])) y1 = max(0, int(rect["y"])) x2 = min(w, int(rect["x"] + rect["width"])) y2 = min(h, int(rect["y"] + rect["height"])) draw.rectangle([x1, y1, x2, y2], fill=255) mask.save(output_mask_path) return {"x1": x1, "y1": y1, "x2": x2, "y2": y2} def stream_training_logs(process): """Read training subprocess output in background thread.""" global training_state try: for line in iter(process.stdout.readline, ""): line_stripped = line.rstrip() if line_stripped: with training_lock: training_state["log_lines"].append(line_stripped) training_state["log_lines"] = training_state["log_lines"][-2000:] process.stdout.close() process.wait() with training_lock: training_state["running"] = False training_state["completed"] = (process.returncode == 0) if process.returncode == 0: object_class = training_state.get("object_class") if object_class: sample_img = None good_dir = TRAINING_DATA_DIR / object_class / "test" / "good" if good_dir.exists(): for f in sorted(good_dir.iterdir()): if f.is_file() and allowed_file(f.name): sample_img = f break if not sample_img: prod_dir = TRAINING_DATA_DIR / object_class if prod_dir.exists(): for f in sorted(prod_dir.rglob("*")): if f.is_file() and allowed_file(f.name) and "mask" not in f.name.lower(): sample_img = f break if sample_img: save_to_trained_gallery(str(sample_img), object_class) else: print(f"[Gallery Warning] No clean sample image found for '{object_class}'.") else: training_state["error"] = f"Training exited with code {process.returncode}" except Exception as e: with training_lock: training_state["running"] = False training_state["error"] = f"Log thread crashed: {str(e)}" def save_to_trained_gallery(clean_image_path: str, product_name: str): """Copies 1 clean image to the gallery named after the product for DINOv2 matching.""" try: for existing_file in TRAINED_GALLERY_DIR.iterdir(): if existing_file.is_file() and existing_file.stem == product_name: print(f"[Gallery Info] Reference image for '{product_name}' already exists in gallery. Skipping copy.") return True src = Path(clean_image_path) if not src.exists(): print(f"[Gallery Error] Clean image not found at: {clean_image_path}") return False ext = src.suffix.lower() if ext not in ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff']: ext = '.png' dst = TRAINED_GALLERY_DIR / f"{product_name}{ext}" shutil.copy2(src, dst) print(f"[Gallery Success] Added reference image to gallery: {dst}") return True except Exception as e: print(f"[Gallery Error] Failed to save image to gallery: {e}") return False # ============================================================================= # Routes - Pages # ============================================================================= @app.route("/") def index(): return render_template("index.html") # ============================================================================= # Routes - DefectDiffu Config # ============================================================================= @app.route("/api/defectdiffu/config", methods=["GET"]) def get_defectdiffu_config(): """Get saved DefectDiffu model paths.""" try: with open(DEFECTDIFFU_CONFIG_PATH, "r") as f: config = json.load(f) return jsonify({"success": True, "config": config}) except Exception as e: return jsonify({"success": False, "error": str(e)}) @app.route("/api/defectdiffu/config", methods=["POST"]) def save_defectdiffu_config(): """Save DefectDiffu model paths.""" data = request.get_json() config = { "ckpt_path": data.get("ckpt_path", "").strip(), "vae_path": data.get("vae_path", "").strip(), "vlm_model": data.get("vlm_model", "gemma3:12b").strip() } try: with open(DEFECTDIFFU_CONFIG_PATH, "w") as f: json.dump(config, f, indent=2) return jsonify({"success": True, "message": "DefectDiffu config saved."}) except Exception as e: return jsonify({"success": False, "error": str(e)}) # ============================================================================= # Routes - Training Data Upload & Mask (DefectFill) # ============================================================================= @app.route("/api/training/good-images/count", methods=["GET"]) def get_good_image_count(): """Get count of existing good images for a product.""" object_class = request.args.get("object_class", "").strip() if not object_class: return jsonify({"success": False, "error": "Object class required."}) good_dir = TRAINING_DATA_DIR / object_class / "test" / "good" count = 0 if good_dir.exists(): count = len([f for f in good_dir.iterdir() if f.is_file() and allowed_file(f.name)]) return jsonify({ "success": True, "count": count, "path": str(good_dir.relative_to(DEFECTFILL_ROOT)) }) @app.route("/api/training/upload-good-images", methods=["POST"]) def upload_good_images(): """Upload good (non-defective) images to temp staging area.""" object_class = request.form.get("object_class", "").strip() if not object_class: return jsonify({"success": False, "error": "Product name is required."}) if "images" not in request.files: return jsonify({"success": False, "error": "No images provided."}) files = request.files.getlist("images") target_dir = TEMP_IMAGE_DIR / object_class / "test" / "good" if target_dir.exists(): shutil.rmtree(target_dir) target_dir.mkdir(parents=True, exist_ok=True) saved = [] for file in files: if file and allowed_file(file.filename): safe_name = Path(file.filename).name save_path = target_dir / safe_name counter = 1 while save_path.exists(): stem = Path(file.filename).stem suffix = Path(file.filename).suffix save_path = target_dir / f"{stem}_{counter:02d}{suffix}" counter += 1 file.save(save_path) img = cv2.imread(str(save_path)) if img is not None: resized = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LANCZOS4) cv2.imwrite(str(save_path), resized) rel = save_path.relative_to(TEMP_IMAGE_DIR).as_posix() saved.append({ "filename": save_path.name, "path": rel, "url": f"/api/temp/image/{rel}" }) return jsonify({ "success": True, "saved": saved, "count": len(saved), "total_in_folder": len([f for f in target_dir.iterdir() if f.is_file() and allowed_file(f.name)]), "target_dir": str(target_dir.relative_to(DEFECTFILL_ROOT).as_posix()) }) @app.route("/api/training/good-image/") def serve_good_image(subpath): """Serve uploaded good images.""" return send_from_directory(TRAINING_DATA_DIR, subpath) @app.route("/api/training/upload-images", methods=["POST"]) def upload_training_images(): """Upload defect images to temp staging area.""" object_class = request.form.get("object_class", "").strip() defect_type = request.form.get("defect_type", "").strip() if not object_class or not defect_type: return jsonify({"success": False, "error": "Product name and defect type are required."}) if "images" not in request.files: return jsonify({"success": False, "error": "No images provided."}) files = request.files.getlist("images") if len(files) < 1: return jsonify({"success": False, "error": f"At least 1 defect images required. Received {len(files)}."}) target_dir = TEMP_IMAGE_DIR / object_class / "train" / "defective" / defect_type if target_dir.exists(): shutil.rmtree(target_dir) target_dir.mkdir(parents=True, exist_ok=True) saved = [] for file in files: if file and allowed_file(file.filename): filename = file.filename safe_name = Path(filename).name save_path = target_dir / safe_name counter = 1 while save_path.exists(): stem = Path(filename).stem suffix = Path(filename).suffix save_path = target_dir / f"{stem}_{counter:02d}{suffix}" counter += 1 file.save(save_path) rel = save_path.relative_to(TEMP_IMAGE_DIR).as_posix() saved.append({ "filename": save_path.name, "path": rel, "url": f"/api/temp/image/{rel}" }) return jsonify({ "success": True, "saved": saved, "count": len(saved), "target_dir": str(target_dir.relative_to(DEFECTFILL_ROOT).as_posix()) }) @app.route("/api/training/image/") def serve_training_image(subpath): """Serve uploaded training images.""" return send_from_directory(TRAINING_DATA_DIR, subpath) @app.route("/api/temp/image/") def serve_temp_image(subpath): """Serve images from the temp staging directory.""" return send_from_directory(TEMP_IMAGE_DIR, subpath) @app.route("/api/training/save-mask", methods=["POST"]) def save_mask(): """Save user-drawn rectangle mask for a training image.""" data = request.get_json() image_path = data.get("image_path", "") rect = data.get("rect", {}) if not image_path or not rect: return jsonify({"success": False, "error": "Missing image_path or rect data."}) full_path = TRAINING_DATA_DIR / image_path if not full_path.exists(): return jsonify({"success": False, "error": "Image not found."}) rel_path = Path(image_path) new_parts = [] for part in rel_path.parts: if part == "defective": new_parts.append("defective_masks") else: new_parts.append(part) mask_path = TRAINING_DATA_DIR / Path(*new_parts) mask_path.parent.mkdir(parents=True, exist_ok=True) coords = generate_mask_from_rect(full_path, rect, mask_path) return jsonify({ "success": True, "mask_path": str(mask_path.relative_to(TRAINING_DATA_DIR)), "coords": coords }) @app.route("/api/training/align", methods=["POST"]) def align_images(): """Align all temp images for a product to a reference image, then move to final.""" if not CV2_AVAILABLE: return jsonify({"success": False, "error": "OpenCV not installed. Cannot align images."}) object_class = request.form.get("object_class", "").strip() reference = request.files.get("reference") if not object_class or not reference: return jsonify({"success": False, "error": "Object class and reference image are required."}) ref_dir = TEMP_IMAGE_DIR / object_class / "reference" ref_dir.mkdir(parents=True, exist_ok=True) ref_path = ref_dir / "reference.png" reference.save(ref_path) obj_temp_dir = TEMP_IMAGE_DIR / object_class aligned = [] errors = [] for img_path in sorted(obj_temp_dir.rglob("*")): if not img_path.is_file(): continue if "reference" in str(img_path.relative_to(obj_temp_dir)).split(os.sep): continue if not allowed_file(img_path.name): continue rel = img_path.relative_to(obj_temp_dir) final_path = TRAINING_DATA_DIR / object_class / rel final_path.parent.mkdir(parents=True, exist_ok=True) try: align_image_to_reference(str(ref_path), str(img_path), str(final_path)) img_arr = cv2.imread(str(final_path)) if img_arr is not None: img_arr_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4) cv2.imwrite(str(final_path), img_arr_resized) aligned.append({ "filename": final_path.name, "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()), "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}" }) except Exception as e: errors.append({"file": str(rel), "error": str(e)}) shutil.copy2(str(img_path), str(final_path)) img_arr = cv2.imread(str(final_path)) if img_arr is not None: img_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4) cv2.imwrite(str(final_path), img_resized) aligned.append({ "filename": final_path.name, "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()), "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}" }) return jsonify({ "success": True, "aligned_count": len(aligned), "error_count": len(errors), "errors": errors, "images": aligned }) @app.route("/api/training/skip-align", methods=["POST"]) def skip_align(): """Copy all temp images directly to final location without alignment.""" data = request.get_json() object_class = data.get("object_class", "").strip() if not object_class: return jsonify({"success": False, "error": "Object class required."}) obj_temp_dir = TEMP_IMAGE_DIR / object_class moved = [] if obj_temp_dir.exists(): for img_path in sorted(obj_temp_dir.rglob("*")): if not img_path.is_file(): continue if "reference" in str(img_path.relative_to(obj_temp_dir)).split(os.sep): continue rel = img_path.relative_to(obj_temp_dir) final_path = TRAINING_DATA_DIR / object_class / rel final_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(str(img_path), str(final_path)) img_arr = cv2.imread(str(final_path)) if img_arr is not None: img_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4) cv2.imwrite(str(final_path), img_resized) moved.append({ "filename": final_path.name, "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()), "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}" }) return jsonify({"success": True, "moved_count": len(moved), "images": moved}) # ============================================================================= # Routes - Domain Instructions # ============================================================================= @app.route("/api/domain-instructions", methods=["GET"]) def get_domain_instructions(): """Get saved domain instructions.""" try: with open(DOMAIN_CONFIG_PATH, "r") as f: config = json.load(f) return jsonify({"success": True, "config": config}) except Exception as e: return jsonify({"success": False, "error": str(e)}) @app.route("/api/domain-instructions", methods=["POST"]) def save_domain_instructions(): """Save domain-specific instructions for prompts.""" data = request.get_json() object_class = data.get("object_class", "default").strip() instructions = data.get("instructions", "").strip() try: with open(DOMAIN_CONFIG_PATH, "r") as f: config = json.load(f) if "active" not in config: config["active"] = {} # ========================================================== # NEW APPEND LOGIC STARTS HERE # ========================================================== existing_instructions = config["active"].get(object_class, "") if instructions.strip(): # If the user wrote something in the textarea, append it to the old instructions if existing_instructions.strip(): # Use double newlines to clearly separate the old and new contexts config["active"][object_class] = existing_instructions.strip() + "\n\n" + instructions.strip() else: # If there were no old instructions, just use the new ones config["active"][object_class] = instructions.strip() else: # If the user left the textarea EMPTY, we do NOT erase old instructions. # We simply keep the existing instructions untouched. pass # ========================================================== # NEW APPEND LOGIC ENDS HERE # ========================================================== with open(DOMAIN_CONFIG_PATH, "w") as f: json.dump(config, f, indent=2) return jsonify({"success": True, "message": f"Instructions saved for '{object_class}'."}) except Exception as e: return jsonify({"success": False, "error": str(e)}) # ============================================================================= # Routes - Training Execution (DefectFill) # ============================================================================= @app.route("/api/check-product-exists", methods=["GET"]) def check_product_exists(): """Check if a product (object_class) already exists in training data.""" object_class = request.args.get("object_class", "").strip() if not object_class: return jsonify({"success": False, "error": "Product name is required."}) product_dir = TRAINING_DATA_DIR / object_class exists = product_dir.exists() and product_dir.is_dir() existing_defect_types = [] sample_image_url = None if exists: defective_dir = product_dir / "defective" if defective_dir.exists(): for defect_dir in sorted(defective_dir.iterdir()): if defect_dir.is_dir(): existing_defect_types.append(defect_dir.name) if not sample_image_url: for img_file in sorted(defect_dir.iterdir()): if allowed_file(img_file.name): rel_path = img_file.relative_to(BASE_DIR) sample_image_url = f"/api/training/image/{rel_path}" break return jsonify({ "success": True, "exists": exists, "existing_defect_types": existing_defect_types, "sample_image_url": sample_image_url }) @app.route("/api/training/start", methods=["POST"]) def start_training(): """Start few-shot training with train.py.""" global training_state data = request.get_json() object_class = data.get("object_class", "").strip() defect_type = data.get("defect_type", "").strip() lora_rank = data.get("lora_rank", 8) lora_alpha = data.get("lora_alpha", 16) max_steps = data.get("max_train_steps", 1500) batch_size = data.get("batch_size", 2) gradient_accum = data.get("gradient_accumulation_steps", 2) lambda_defect = data.get("lambda_defect", 0.5) lambda_obj = data.get("lambda_obj", 0.2) lambda_attn = data.get("lambda_attn", 0.05) alpha = data.get("alpha", 0.3) if not object_class or not defect_type: return jsonify({"success": False, "error": "Object class and defect type are required."}) with training_lock: if training_state["running"]: return jsonify({"success": False, "error": "Training is already in progress."}) output_dir = CHECKPOINT_DIR / object_class / defect_type output_dir.mkdir(parents=True, exist_ok=True) train_script = TRAIN_PY if not train_script.exists(): train_script = BASE_DIR.parent / "train.py" cmd = [ sys.executable, '-u', str(train_script), "--data_dir", str(TRAINING_DATA_DIR), "--object_class", object_class, "--defect_type", defect_type, "--output_dir", str(output_dir), "--lora_rank", str(lora_rank), "--lora_alpha", str(lora_alpha), "--max_train_steps", str(max_steps), "--batch_size", str(batch_size), "--gradient_accumulation_steps", str(gradient_accum), "--lambda_defect", str(lambda_defect), "--lambda_obj", str(lambda_obj), "--lambda_attn", str(lambda_attn), "--alpha", str(alpha), "--save_steps", str(max(500, max_steps // 3)), "--lr_warmup_steps", str(min(100, max_steps // 10)), "--dilate_mask", "False", "--seed", str(int(time.time() * 1000) % (2**31)) ] try: with training_lock: training_state["running"] = True training_state["object_class"] = object_class training_state["defect_type"] = defect_type training_state["output_dir"] = str(output_dir) training_state["log_lines"] = [] training_state["start_time"] = datetime.now().isoformat() training_state["completed"] = False training_state["error"] = None process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, universal_newlines=True, cwd=str(DEFECTFILL_ROOT) ) with training_lock: training_state["process"] = process log_thread = threading.Thread(target=stream_training_logs, args=(process,)) log_thread.daemon = True log_thread.start() return jsonify({ "success": True, "message": "Training started.", "command": " ".join(cmd), "output_dir": str(output_dir.relative_to(DEFECTFILL_ROOT)) }) except Exception as e: with training_lock: training_state["running"] = False training_state["error"] = str(e) return jsonify({"success": False, "error": str(e)}) @app.route("/api/training/status", methods=["GET"]) def training_status(): """Get current training status and logs.""" with training_lock: state = dict(training_state) state["process"] = None return jsonify({"success": True, "status": state}) @app.route("/api/training/stop", methods=["POST"]) def stop_training(): """Stop running training.""" global training_state with training_lock: if training_state["process"] and training_state["process"].poll() is None: training_state["process"].terminate() training_state["running"] = False training_state["error"] = "Training stopped by user." return jsonify({"success": True, "message": "Training terminated."}) return jsonify({"success": False, "error": "No training is running."}) # ============================================================================= # Routes - Checkpoints & Generation (DefectFill) # ============================================================================= @app.route("/api/detect-similar-product", methods=["POST"]) def detect_similar_product(): """Endpoint to identify if a clean product image has already been trained.""" if "image" not in request.files: return jsonify({"success": False, "error": "No clean product image uploaded."}) file = request.files["image"] temp_path = UPLOAD_FOLDER / f"temp_detect_{uuid.uuid4().hex}{Path(file.filename).suffix}" try: file.save(temp_path) threshold = float(request.form.get("threshold", 0.85)) result = check_product_training_status( target_image_path=str(temp_path), gallery_dir=TRAINED_GALLERY_DIR, similarity_threshold=threshold ) return jsonify({"success": True, "result": result}) except Exception as e: return jsonify({"success": False, "error": str(e)}) finally: if temp_path.exists(): os.remove(temp_path) @app.route("/api/checkpoints", methods=["GET"]) def list_checkpoints(): """List available trained checkpoints.""" return jsonify({"success": True, "checkpoints": get_checkpoints()}) @app.route("/api/generate", methods=["POST"]) def generate_defect(): """Start DefectFill generation in background and return a job_id for polling.""" object_class = request.form.get("object_class", "").strip() product_desc = request.form.get("product_desc", "").strip() detection_mode = request.form.get("detection_mode", "dots") if not object_class or not product_desc: return jsonify({"success": False, "error": "Object class and product description are required."}) defect_type = request.form.get("defect_type", "").strip() or None num_defects = int(request.form.get("num_defects", 3)) guidance_scale = float(request.form.get("guidance_scale", 8.0)) num_steps = int(request.form.get("num_steps", 50)) image_size = int(request.form.get("image_size", 512)) device = request.form.get("device", "cuda") clean_files = request.files.getlist("clean_images") if not clean_files: return jsonify({"success": False, "error": "Clean images are required."}) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") allowed_exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp'} clean_paths = [] for f in clean_files: if not f or not f.filename: continue ext = os.path.splitext(f.filename)[1].lower() if ext not in allowed_exts: continue safe_name = os.path.basename(f.filename) clean_filename = f"clean_{timestamp}_{safe_name}" clean_path = UPLOAD_FOLDER / clean_filename f.save(clean_path) clean_paths.append(str(clean_path)) if not clean_paths: return jsonify({"success": False, "error": "No valid image files found in selected folder."}) ckpt_dir = CHECKPOINT_DIR valid_types = get_checkpoints().get(object_class, []) if defect_type and defect_type not in valid_types: return jsonify({ "success": False, "error": f"Defect type '{defect_type}' not found for '{object_class}'. Available: {valid_types}" }) valid_object_classes = list(get_checkpoints().keys()) valid_defect_types = get_checkpoints() gen_output_dir = OUTPUT_DIR / f"gen_{timestamp}" gen_output_dir.mkdir(parents=True, exist_ok=True) try: domain_hint = "" if DOMAIN_CONFIG_PATH.exists(): with open(DOMAIN_CONFIG_PATH, "r") as f: dconfig = json.load(f) domain_hint = dconfig.get("active", {}).get(object_class, "") # If detection_mode == "auto", do nothing - let the VLM infer entirely from the product description and domain context. if detection_mode == "lines": product_desc += " Focus defect placement on metallic pins, leads, or legs." elif detection_mode == "dots": product_desc += " Focus defect placement on solder ball array, BGA grid, or rounded pads." elif detection_mode == "single_rounded": product_desc += " Focus defect placement on a single, isolated rounded feature such as a mounting hole, circular pad, individual via, or isolated dot." # if domain_hint: # product_desc += f" Domain context: {domain_hint}" job_params = { 'device': device, 'gen_output_dir': gen_output_dir, 'ckpt_dir': ckpt_dir, 'object_class': object_class, 'defect_type': defect_type, 'valid_object_classes': valid_object_classes, 'valid_defect_types': valid_defect_types, 'image_size': image_size, 'num_steps': num_steps, 'guidance_scale': guidance_scale, 'product_desc': product_desc, 'domain_hint': domain_hint, 'clean_paths': clean_paths, 'num_defects': num_defects, } job_id = str(uuid.uuid4()) with jobs_lock: generation_jobs[job_id] = { 'status': 'starting', 'step': 0, 'total_steps': 5, 'step_text': 'Initializing...', 'defect_current': 0, 'defect_total': 0, 'image_current': 0, 'image_total': 0, 'result': None, 'error': None, 'has_error_trace': False } thread = threading.Thread(target=run_generation_job, args=(job_id, job_params, 'defectfill'), daemon=True) thread.start() return jsonify({'success': True, 'job_id': job_id}) except Exception as e: traceback.print_exc() return jsonify({"success": False, "error": str(e)}) # ============================================================================= # Routes - DefectDiffu Generation # ============================================================================= @app.route("/api/defectdiffu/generate", methods=["POST"]) def generate_defectdiffu(): """Start DefectDiffu generation in background and return a job_id for polling.""" product_desc = request.form.get("product_desc", "").strip() defect_type = request.form.get("defect_type", "").strip() or None num_defects = int(request.form.get("num_defects", 3)) image_size = int(request.form.get("image_size", 512)) num_steps = int(request.form.get("num_steps", 50)) device = request.form.get("device", "cuda") if not product_desc: return jsonify({"success": False, "error": "Product description is required."}) # Load model config try: with open(DEFECTDIFFU_CONFIG_PATH, "r") as f: dd_config = json.load(f) except Exception as e: return jsonify({"success": False, "error": f"Failed to load DefectDiffu config: {e}"}) ckpt_path = dd_config.get("ckpt_path", "") vae_path = dd_config.get("vae_path", "") vlm_model = dd_config.get("vlm_model", "gemma3:12b") if not ckpt_path or not vae_path: return jsonify({"success": False, "error": "DefectDiffu model paths not configured. Please set them in the UI first."}) if not Path(ckpt_path).exists(): return jsonify({"success": False, "error": f"Checkpoint not found: {ckpt_path}"}) if not Path(vae_path).exists(): return jsonify({"success": False, "error": f"VAE not found: {vae_path}"}) clean_files = request.files.getlist("clean_images") if not clean_files: return jsonify({"success": False, "error": "Clean images are required."}) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") allowed_exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp'} clean_paths = [] for f in clean_files: if not f or not f.filename: continue ext = os.path.splitext(f.filename)[1].lower() if ext not in allowed_exts: continue safe_name = os.path.basename(f.filename) clean_filename = f"clean_dd_{timestamp}_{safe_name}" clean_path = UPLOAD_FOLDER / clean_filename f.save(clean_path) clean_paths.append(str(clean_path)) if not clean_paths: return jsonify({"success": False, "error": "No valid image files found."}) gen_output_dir = OUTPUT_DIR / f"defectdiffu_{timestamp}" gen_output_dir.mkdir(parents=True, exist_ok=True) try: job_params = { 'device': device, 'gen_output_dir': gen_output_dir, 'defectdiffu_ckpt': ckpt_path, 'vae_path': vae_path, 'vlm_model': vlm_model, 'image_size': image_size, 'num_steps': num_steps, 'product_desc': product_desc, 'clean_paths': clean_paths, 'num_defects': num_defects, 'defect_type': defect_type, } job_id = str(uuid.uuid4()) with jobs_lock: generation_jobs[job_id] = { 'status': 'starting', 'step': 0, 'total_steps': 5, 'step_text': 'Initializing DefectDiffu...', 'defect_current': 0, 'defect_total': 0, 'image_current': 0, 'image_total': 0, 'result': None, 'error': None, 'has_error_trace': False } thread = threading.Thread(target=run_generation_job, args=(job_id, job_params, 'defectdiffu'), daemon=True) thread.start() return jsonify({'success': True, 'job_id': job_id}) except Exception as e: traceback.print_exc() return jsonify({"success": False, "error": str(e)}) @app.route('/api/generation-status/', methods=['GET']) def generation_status(job_id): with jobs_lock: job = generation_jobs.get(job_id) if not job: return jsonify({'success': False, 'error': 'Job not found'}), 404 return jsonify({'success': True, 'job': job}) @app.route("/api/output/image/") def serve_output_image(subpath): """Serve generated output images.""" return send_from_directory(BASE_DIR, subpath) # ============================================================================= # Routes - System Info # ============================================================================= @app.route("/api/system/info", methods=["GET"]) def system_info(): """Get system information.""" import torch info = { "cuda_available": torch.cuda.is_available(), "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0, "cuda_device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, "project_root": str(PROJECT_ROOT), "defectfill_root": str(DEFECTFILL_ROOT), "defectdiffu_root": str(DEFECTDIFFU_ROOT), "checkpoint_dir": str(CHECKPOINT_DIR), "training_data_dir": str(TRAINING_DATA_DIR) } return jsonify({"success": True, "info": info}) # ============================================================================= # Main # ============================================================================= if __name__ == "__main__": print("=" * 60) print(" DefectFill & DefectDiffu Unified Web UI") print("=" * 60) print(f" Project root: {PROJECT_ROOT}") print(f" DefectFill root: {DEFECTFILL_ROOT}") print(f" DefectDiffu root: {DEFECTDIFFU_ROOT}") print(f" Checkpoints: {CHECKPOINT_DIR}") print(f" Training data: {TRAINING_DATA_DIR}") print(f" Outputs: {OUTPUT_DIR}") print("-" * 60) print(" Open http://127.0.0.1:5000 in your browser") print("=" * 60) app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)