""" FloorplanVLM GRPO Training (Stage 2) - Run after SFT Loads the SFT model and applies geometric reward-based RL. Reward: R = 0.1·R_val + 0.5·R_ext + α·0.4·R_int (FloorplanVLM Eq. 9) Usage: pip install shapely # + same deps as SFT script python train_floorplan_grpo.py Requires: SFT model already pushed to HUB (default: manitocross/floorplan-vlm-sft) """ import os, json, re, math, torch, numpy as np from PIL import Image, ImageDraw from datasets import Dataset from transformers import AutoProcessor, TrainerCallback from trl import GRPOTrainer, GRPOConfig from peft import LoraConfig from shapely.geometry import Polygon, LineString from shapely.ops import unary_union # ══════════════════════════════════════════════════════════════════════════════ # CONFIGURATION # ══════════════════════════════════════════════════════════════════════════════ SFT_MODEL_ID = "manitocross/floorplan-vlm-sft" # your SFT model from Stage 1 HUB_MODEL_ID = "manitocross/floorplan-vlm-grpo" OUTPUT_DIR = "./floorplan-vlm-grpo" # Uses same CubiCasa5K data directory as SFT script DATA_DIR = "./cubicasa_data" NUM_EPOCHS = 1 BATCH_SIZE = 1 GRAD_ACCUM = 4 LEARNING_RATE = 1e-6 NUM_GENERATIONS = 4 # G: completions per prompt MAX_COMPLETION_LENGTH = 4096 KL_COEF = 0.01 PUSH_TO_HUB = True # ══════════════════════════════════════════════════════════════════════════════ SYSTEM_PROMPT = "You are a floor plan vectorization expert. Extract wall, door, window geometry from floor plan images into structured JSON. Output ONLY valid JSON." USER_PROMPT = "Vectorize this floor plan into structured JSON with all walls, doors, windows, and rooms." # ── Geometric Helpers ──────────────────────────────────────────────────────── def extract_json(text): if isinstance(text, list): text = text[0].get("content", "") if text else "" text = text.strip() try: return json.loads(text) except: pass m = re.search(r'\{[\s\S]*\}', text) if m: try: return json.loads(m.group()) except: pass return None def walls_to_polygon(walls): if not walls or len(walls) < 3: return None try: polys = [] for w in walls: s, e = w.get('start',[0,0]), w.get('end',[0,0]) t = max(w.get('thickness',10), 1) polys.append(LineString([s, e]).buffer(t/2, cap_style=2)) combined = unary_union(polys) return combined.convex_hull if not combined.is_empty else None except: return None def poly_iou(p1, p2): if p1 is None or p2 is None: return 0.0 try: if not p1.is_valid: p1 = p1.buffer(0) if not p2.is_valid: p2 = p2.buffer(0) inter = p1.intersection(p2).area union = p1.union(p2).area return inter / union if union > 0 else 0.0 except: return 0.0 # ── Reward Function ────────────────────────────────────────────────────────── def floorplan_reward(completions, **kwargs): """Combined: 0.1·R_val + 0.5·R_ext + α·0.4·R_int""" gt_jsons = kwargs.get("json_gt", []) rewards = [] for c, gt_str in zip(completions, gt_jsons): text = c[0]["content"] if isinstance(c, list) else c pred = extract_json(text) if pred is None: rewards.append(0.0); continue try: gt = json.loads(gt_str) if isinstance(gt_str, str) else gt_str except: rewards.append(0.0); continue # R_val r_val = 0.0 has_walls = "walls" in pred and isinstance(pred["walls"], list) and len(pred["walls"]) > 0 if has_walls: valid = sum(1 for w in pred["walls"] if all(k in w for k in ["id","start","end","thickness"]) and isinstance(w.get("start"), list) and len(w.get("start",[])) == 2) r_val = 0.3 + 0.5 * (valid / max(len(pred["walls"]), 1)) if "rooms" in pred and isinstance(pred["rooms"], list): wids = {w.get("id") for w in pred["walls"]} vr = sum(1 for r in pred["rooms"] if "label" in r and "walls" in r and all(wid in wids for wid in r.get("walls",[]))) r_val += 0.2 * (vr / max(len(pred["rooms"]), 1)) # R_ext pred_poly = walls_to_polygon(pred.get("walls", [])) gt_poly = walls_to_polygon(gt.get("walls", [])) r_ext = poly_iou(pred_poly, gt_poly) # Alpha gating (FloorplanVLM Eq. 8) if r_ext < 0.3: alpha = 0.1 elif r_ext < 0.7: alpha = 0.1 + 0.9 * (r_ext - 0.3) / 0.4 else: alpha = 1.0 # R_int (room label overlap) r_int = 0.0 pred_rooms = pred.get("rooms", []) gt_rooms = gt.get("rooms", []) if pred_rooms and gt_rooms: pl = set(r.get("label","") for r in pred_rooms) gl = set(r.get("label","") for r in gt_rooms) overlap = len(pl & gl) total = len(pl | gl) r_int = overlap / total if total > 0 else 0.0 total = 0.1 * min(r_val, 1.0) + 0.5 * r_ext + alpha * 0.4 * r_int rewards.append(float(total)) return rewards # ── Dataset (reuses SFT data dir) ─────────────────────────────────────────── # You'll need to adapt this to load your data (same CubiCasa5K directory). # For GRPO format: prompt-only (no completion), plus metadata for reward. def build_grpo_dataset(data_dir, max_samples=100): """Build GRPO dataset — loads pre-converted JSON annotations.""" # Look for the annotation file from SFT stage ann_path = os.path.join(data_dir, "annotations.json") if not os.path.exists(ann_path): print(f" No pre-converted annotations at {ann_path}") print(f" Run train_floorplan_vlm.py (SFT) first to generate data.") print(f" Creating synthetic GRPO dataset as fallback...") return create_synthetic_grpo(max_samples or 20) with open(ann_path) as f: annotations = json.load(f) if max_samples: annotations = annotations[:max_samples] records = [] for ann in annotations: img_path = ann.get("image_path") if not img_path or not os.path.exists(img_path): continue img = Image.open(img_path).convert("RGB") records.append({ "prompt": [ {"role":"system","content":[{"type":"text","text":SYSTEM_PROMPT}]}, {"role":"user","content":[{"type":"image"},{"type":"text","text":USER_PROMPT}]}, ], "images": [img], "json_gt": ann["json_annotation"], }) print(f" GRPO dataset: {len(records)} samples") return Dataset.from_list(records) def create_synthetic_grpo(n=20): records = [] for i in range(n): size = 256 img = Image.new('RGB', (size,size), 'white') d = ImageDraw.Draw(img) m = 30+i*3; wt=6; s=1024.0/size; mid=size//2+i*2 d.rectangle([m,m,size-m,size-m], outline='black', width=wt) d.line([(m,mid),(size-m,mid)], fill='black', width=wt) jd = {"walls":[ {"id":"wall_1","start":[round(m*s),round(m*s)],"end":[round((size-m)*s),round(m*s)],"thickness":round(wt*s),"curvature":0,"openings":[]}, {"id":"wall_2","start":[round((size-m)*s),round(m*s)],"end":[round((size-m)*s),round((size-m)*s)],"thickness":round(wt*s),"curvature":0,"openings":[]}, {"id":"wall_3","start":[round((size-m)*s),round((size-m)*s)],"end":[round(m*s),round((size-m)*s)],"thickness":round(wt*s),"curvature":0,"openings":[]}, {"id":"wall_4","start":[round(m*s),round((size-m)*s)],"end":[round(m*s),round(m*s)],"thickness":round(wt*s),"curvature":0,"openings":[]}, {"id":"wall_5","start":[round(m*s),round(mid*s)],"end":[round((size-m)*s),round(mid*s)],"thickness":round(wt*s),"curvature":0,"openings":[]}, ],"rooms":[ {"label":"bedroom","walls":["wall_1","wall_2","wall_5","wall_4"]}, {"label":"living_room","walls":["wall_5","wall_2","wall_3","wall_4"]}, ]} records.append({ "prompt":[ {"role":"system","content":[{"type":"text","text":SYSTEM_PROMPT}]}, {"role":"user","content":[{"type":"image"},{"type":"text","text":USER_PROMPT}]}, ], "images":[img], "json_gt": json.dumps(jd, separators=(',',':')), }) return Dataset.from_list(records) # ── Main ───────────────────────────────────────────────────────────────────── def main(): use_gpu = torch.cuda.is_available() print("="*64) print(f" FloorplanVLM GRPO Training ({'GPU' if use_gpu else 'CPU'})") print(f" SFT Model : {SFT_MODEL_ID}") print(f" Output : {HUB_MODEL_ID}") print(f" Generations: {NUM_GENERATIONS}, KL: {KL_COEF}") print("="*64) dataset = build_grpo_dataset(DATA_DIR) peft_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", ) grpo_config = GRPOConfig( output_dir=OUTPUT_DIR, num_train_epochs=NUM_EPOCHS, per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRAD_ACCUM, learning_rate=LEARNING_RATE, bf16=use_gpu, fp16=False, gradient_checkpointing=True, logging_steps=5, logging_first_step=True, logging_strategy="steps", disable_tqdm=True, save_steps=200, save_total_limit=2, num_generations=NUM_GENERATIONS, max_prompt_length=512, max_completion_length=MAX_COMPLETION_LENGTH, scale_rewards=True, beta=KL_COEF, temperature=0.7, push_to_hub=PUSH_TO_HUB, hub_model_id=HUB_MODEL_ID if PUSH_TO_HUB else None, report_to="none", ) trainer = GRPOTrainer( model=SFT_MODEL_ID, reward_funcs=[floorplan_reward], args=grpo_config, train_dataset=dataset, peft_config=peft_config, ) print("\nStarting GRPO training...") trainer.train() trainer.save_model(OUTPUT_DIR) if PUSH_TO_HUB: try: trainer.push_to_hub() print(f"\n✅ Model: https://huggingface.co/{HUB_MODEL_ID}") except Exception as e: print(f"Push failed: {e}") from huggingface_hub import HfApi HfApi().create_repo(HUB_MODEL_ID, exist_ok=True) HfApi().upload_folder(folder_path=OUTPUT_DIR, repo_id=HUB_MODEL_ID) print(f"\n✅ Done! Local: {OUTPUT_DIR}/") if __name__ == "__main__": main()