import copy import json import math import os import sys from typing import Iterable import cv2 import numpy as np import torch from shapely.geometry import Polygon import util.misc as utils # Add evaluations folder relative to this file's location sys.path.append(os.path.join(os.path.dirname(__file__), "evaluations")) from rplan_eval.Evaluator import Evaluator_RPlan from s3d_floorplan_eval.DataRW.S3DRW import S3DRW from s3d_floorplan_eval.DataRW.wrong_annotatios import wrong_s3d_annotations_list from s3d_floorplan_eval.Evaluator.Evaluator import Evaluator from s3d_floorplan_eval.options import MCSSOptions from datasets import get_dataset_class_labels from util.eval_utils import compute_f1 from util.plot_utils import ( concat_floorplan_maps, plot_density_map, plot_floorplan_with_regions, plot_semantic_rich_floorplan_opencv, sort_polygons_by_matching, ) from util.poly_ops import pad_gt_polys options = MCSSOptions() opts = options.parse() def train_one_epoch( model: torch.nn.Module, criterion: torch.nn.Module, data_loader: Iterable, optimizer: torch.optim.Optimizer, device: torch.device, epoch: int, max_norm: float = 0, poly2seq: bool = False, ema_model=None, **kwargs, ): """ Trains the model for one epoch using the provided data loader, criterion, and optimizer. This function iterates over the data loader, computes losses, performs backpropagation, applies gradient clipping if specified, updates the model parameters, and optionally updates an EMA model. It logs various metrics including loss, learning rate, and gradient norm. Args: model (torch.nn.Module): The neural network model to be trained. criterion (torch.nn.Module): The loss criterion used to compute the loss. data_loader (Iterable): An iterable data loader yielding batches of inputs and extras. optimizer (torch.optim.Optimizer): The optimizer used to update model parameters. device (torch.device): The device (CPU or GPU) on which to perform computations. epoch (int): The current epoch number, used for logging. max_norm (float, optional): Maximum norm for gradient clipping. If 0, no clipping is applied. Defaults to 0. poly2seq (bool, optional): If True, uses batched_extras as room_targets and passes them to the model. Defaults to False. ema_model (optional): Exponential moving average model to update, if provided. Defaults to None. **kwargs: Additional keyword arguments, such as 'drop_rate' for padding ground truth polygons. Returns: dict: A dictionary containing the global averages of logged metrics (e.g., loss, lr, grad_norm). """ model.train() criterion.train() metric_logger = utils.MetricLogger(delimiter=" ") metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}")) metric_logger.add_meter("grad_norm", utils.SmoothedValue(window_size=1, fmt="{value:.2f}")) header = "Epoch: [{}]".format(epoch) print_freq = 10 model_obj = model if not hasattr(model, "module") else model.module for batched_inputs, batched_extras in metric_logger.log_every(data_loader, print_freq, header): samples = [x["image"].to(device) for x in batched_inputs] gt_instances = [x["instances"].to(device) for x in batched_inputs] if not poly2seq: room_targets = pad_gt_polys( gt_instances, model_obj.num_queries_per_poly, samples[0].shape[1], drop_rate=kwargs.get("drop_rate", 0.0), device=device, ) outputs = model(samples) else: for key in batched_extras.keys(): batched_extras[key] = batched_extras[key].to(device) room_targets = batched_extras outputs = model(samples, batched_extras) loss_dict = criterion(outputs, room_targets) weight_dict = criterion.weight_dict losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict) loss_dict_unscaled = {f"{k}_unscaled": v for k, v in loss_dict.items()} loss_dict_scaled = {k: v * weight_dict[k] for k, v in loss_dict.items() if k in weight_dict} losses_reduced_scaled = sum(loss_dict_scaled.values()) loss_value = losses_reduced_scaled.item() if not math.isfinite(loss_value): print("Loss is {}, stopping training".format(loss_value)) print(loss_dict) sys.exit(1) optimizer.zero_grad() losses.backward() if max_norm > 0: grad_total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) else: grad_total_norm = utils.get_total_grad_norm(model.parameters(), max_norm) optimizer.step() if ema_model is not None: utils.update_ema(ema_model, model.module, 0.999) metric_logger.update(loss=loss_value, **loss_dict_scaled, **loss_dict_unscaled) metric_logger.update(lr=optimizer.param_groups[0]["lr"]) metric_logger.update(grad_norm=grad_total_norm) print("Averaged stats:", metric_logger) return {k: meter.global_avg for k, meter in metric_logger.meters.items()} @torch.no_grad() def evaluate( model, criterion, dataset_name, data_loader, device, plot_density=False, output_dir=None, epoch=None, poly2seq: bool = False, add_cls_token=False, per_token_sem_loss=False, wd_as_line=True, ): """ Evaluates the model on a given dataset during training, computing losses and various metrics such as room IoU, precision, recall, corner precision/recall, angle precision/recall, and semantic metrics if applicable. This function supports two evaluation modes: - Non-poly2seq mode (poly2seq=False): RoomFormer evaluation with fg_mask filtering - Poly2seq mode (poly2seq=True): Raster2Seq evaluation with sequence-based predictions Args: model: The neural network model to evaluate. criterion: The loss criterion used for evaluation. dataset_name (str): Name of the dataset (e.g., "stru3d", "cubicasa", "r2g", "waffle"). data_loader: DataLoader providing batches of input data. device: The device (e.g., CPU or GPU) to run the evaluation on. plot_density (bool, optional): If True, plots a density map for the last sample. Defaults to False. output_dir (str, optional): Directory to save output plots if plot_density is True. Defaults to None. epoch (int, optional): Current epoch number, used in plot filenames. Defaults to None. poly2seq (bool, optional): If True, uses sequence-based prediction mode with forward_inference. If False, uses mask-based filtering with fg_mask. Defaults to False. add_cls_token (bool, optional): If True, accounts for class tokens in sequence processing. Only used when poly2seq=True. Defaults to False. per_token_sem_loss (bool, optional): If True, computes semantic loss per token using voting. Only used when poly2seq=True. Defaults to False. wd_as_line (bool, optional): If True, treats windows/doors as lines based on corner count. If False, uses semantic class. Only used when poly2seq=True. Defaults to True. Returns: dict: A dictionary containing averaged evaluation statistics, including losses and metrics like room_iou, room_prec, room_rec, corner_prec, corner_rec, angles_prec, angles_rec, and semantic metrics if applicable. """ model.eval() criterion.eval() if dataset_name == "stru3d": door_window_index = [16, 17] elif dataset_name == "cubicasa": door_window_index = [10, 9] else: door_window_index = [] metric_logger = utils.MetricLogger(delimiter=" ") header = "Test:" model_obj = model if not hasattr(model, "module") else model.module for batched_inputs, batched_extras in metric_logger.log_every(data_loader, 10, header): samples = [x["image"].to(device) for x in batched_inputs] scene_ids = [x["image_id"] for x in batched_inputs] gt_instances = [x["instances"].to(device) for x in batched_inputs] if not poly2seq: room_targets = pad_gt_polys( gt_instances, model_obj.num_queries_per_poly, samples[0].shape[1], drop_rate=0.0, device=device ) outputs = model(samples) else: for key in batched_extras.keys(): batched_extras[key] = batched_extras[key].to(device) room_targets = batched_extras outputs = model(samples, batched_extras) image_size = samples[0].size(2) loss_dict = criterion(outputs, room_targets) weight_dict = criterion.weight_dict if poly2seq: outputs = model_obj.forward_inference(samples) pred_corners = outputs["gen_out"] fg_mask = None else: pred_logits = outputs["pred_logits"] pred_corners = outputs["pred_coords"] fg_mask = torch.sigmoid(pred_logits) > 0.5 # select valid corners bs = outputs["pred_logits"].shape[0] if "pred_room_logits" in outputs: pred_room_logits = outputs["pred_room_logits"] prob = torch.nn.functional.softmax(pred_room_logits, -1) _, pred_room_label = prob[..., :-1].max(-1) # process per scene for i in range(bs): # Prepare ground truth data gt_polys, gt_polys_types = [], [] gt_window_doors = [] gt_window_doors_types = [] for gt_poly, gt_id in zip( gt_instances[i].gt_masks.polygons, gt_instances[i].gt_classes.detach().cpu().tolist() ): gt_poly = gt_poly[0].reshape(-1, 2).astype(np.int32) if gt_id in door_window_index: gt_window_doors.append(gt_poly) gt_window_doors_types.append(gt_id) else: gt_polys.append(gt_poly) gt_polys_types.append(gt_id) # Create evaluator based on dataset if dataset_name == "stru3d": if int(scene_ids[i]) in wrong_s3d_annotations_list: continue curr_opts = copy.deepcopy(opts) curr_opts.scene_id = "scene_0" + str(scene_ids[i]) curr_data_rw = S3DRW(curr_opts, mode="online_eval") evaluator = Evaluator(curr_data_rw, curr_opts, disable_overlapping_filter=poly2seq) elif dataset_name in ["cubicasa", "r2g", "waffle"]: evaluator = Evaluator_RPlan(disable_overlapping_filter=poly2seq, wd_as_line=wd_as_line) print("Running Evaluation for scene %s" % scene_ids[i]) room_polys = [] semantic_rich = "pred_room_logits" in outputs if semantic_rich: room_types = [] window_doors = [] window_doors_types = [] scene_outputs = _process_predictions( pred_corners, i, semantic_rich, poly2seq, fg_mask, image_size, pred_room_label if semantic_rich else None, pred_room_logits if semantic_rich else None, add_cls_token, per_token_sem_loss, wd_as_line, door_window_index, dataset_name, ) room_polys = scene_outputs["room_polys"] room_types = scene_outputs["room_types"] window_doors = scene_outputs["window_doors"] window_doors_types = scene_outputs["window_doors_types"] pred_room_label_per_scene = scene_outputs["pred_room_label_per_scene"] if dataset_name == "stru3d": if not semantic_rich: quant_result_dict_scene = evaluator.evaluate_scene(room_polys=room_polys) else: quant_result_dict_scene = evaluator.evaluate_scene( room_polys=room_polys, room_types=room_types, window_door_lines=window_doors, window_door_lines_types=window_doors_types, ) elif dataset_name in ["cubicasa", "r2g", "waffle"]: if not semantic_rich: quant_result_dict_scene = evaluator.evaluate_scene( room_polys=room_polys, gt_polys=gt_polys, room_types=None, gt_polys_types=gt_polys_types, img_size=(image_size, image_size), ) else: quant_result_dict_scene = evaluator.evaluate_scene( room_polys=room_polys, gt_polys=gt_polys, room_types=room_types, gt_polys_types=gt_polys_types, window_door_lines=window_doors, gt_window_doors_list=gt_window_doors, window_door_lines_types=window_doors_types, gt_window_doors_type_list=gt_window_doors_types, img_size=(image_size, image_size), ) if "room_iou" in quant_result_dict_scene: metric_logger.update(room_iou=quant_result_dict_scene["room_iou"]) metric_logger.update(room_prec=quant_result_dict_scene["room_prec"]) metric_logger.update(room_rec=quant_result_dict_scene["room_rec"]) metric_logger.update(corner_prec=quant_result_dict_scene["corner_prec"]) metric_logger.update(corner_rec=quant_result_dict_scene["corner_rec"]) metric_logger.update(angles_prec=quant_result_dict_scene["angles_prec"]) metric_logger.update(angles_rec=quant_result_dict_scene["angles_rec"]) if semantic_rich: metric_logger.update(room_sem_prec=quant_result_dict_scene["room_sem_prec"]) metric_logger.update(room_sem_rec=quant_result_dict_scene["room_sem_rec"]) metric_logger.update(window_door_prec=quant_result_dict_scene["window_door_prec"]) metric_logger.update(window_door_rec=quant_result_dict_scene["window_door_rec"]) # plot last sample if plot_density and len(room_polys) > 0: pred_room_map = plot_density_map(samples[i], image_size, room_polys, pred_room_label_per_scene) cv2.imwrite(os.path.join(output_dir, "{}_pred_room_map_{}.png".format(scene_ids[i], epoch)), pred_room_map) plot_density = False # only plot once loss_dict_scaled = {k: v * weight_dict[k] for k, v in loss_dict.items() if k in weight_dict} loss_dict_unscaled = {f"{k}_unscaled": v for k, v in loss_dict.items()} metric_logger.update(loss=sum(loss_dict_scaled.values()), **loss_dict_scaled, **loss_dict_unscaled) print("Averaged stats:", metric_logger) stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()} return stats def _process_predictions( pred_corners, i, semantic_rich, poly2seq, fg_mask, image_size, pred_room_label, pred_room_logits, add_cls_token, per_token_sem_loss, wd_as_line, door_window_index, dataset_name, ): """ Processes predictions for room layouts, extracting polygons for rooms, windows, and doors. This function handles two main modes: non-poly2seq (mask-based) and poly2seq (sequence-based). It filters and validates predicted corners based on various conditions, computes room types and window/door types if semantic information is rich, and ensures polygons meet area thresholds. Args: pred_corners (list or tensor): Predicted corner coordinates for scenes or sequences. i (int): Index of the current scene in the batch. semantic_rich (bool): Whether to include semantic information (room types, window/door types). poly2seq (bool): Whether the predictions are in sequence format (poly2seq mode). fg_mask (tensor): Foreground masks for valid corners per room (used in non-poly2seq mode). image_size (int): Size of the image (used for scaling coordinates). pred_room_label (tensor): Predicted room labels for the scene. pred_room_logits (tensor): Logits for room predictions. add_cls_token (int): Number of CLS tokens added in sequences (used in poly2seq mode). per_token_sem_loss (bool): Whether to compute semantic loss per token (affects class aggregation). wd_as_line (bool): Whether to treat windows/doors as lines (affects classification logic). door_window_index (list or dict): Indices or mapping for door/window classes. dataset_name (str): Name of the dataset being processed. Returns: dict: A dictionary containing: - "room_polys" (list of np.ndarray): List of valid room polygons (scaled and rounded). - "room_types" (list or None): List of room types if semantic_rich, else None. - "window_doors" (list or None): List of window/door polygons if semantic_rich, else None. - "window_doors_types" (list or None): List of window/door types if semantic_rich, else None. - "pred_room_label_per_scene" (list): Processed room labels for the scene. """ np_softmax = lambda x: np.exp(x) / np.sum(np.exp(x), axis=-1, keepdims=True) pred_corners_per_scene = pred_corners[i] room_polys = [] if semantic_rich: room_types = [] window_doors = [] window_doors_types = [] pred_room_label_per_scene = pred_room_label[i].cpu().numpy() pred_room_logit_per_scene = pred_room_logits[i].cpu().numpy() # Process predictions based on mode if not poly2seq: fg_mask_per_scene = fg_mask[i] # process per room for j in range(fg_mask_per_scene.shape[0]): fg_mask_per_room = fg_mask_per_scene[j] pred_corners_per_room = pred_corners_per_scene[j] valid_corners_per_room = pred_corners_per_room[fg_mask_per_room] if len(valid_corners_per_room) > 0: corners = (valid_corners_per_room * (image_size - 1)).cpu().numpy() corners = np.around(corners).astype(np.int32) if not semantic_rich: # only regular rooms if len(corners) >= 4 and Polygon(corners).area >= 100: room_polys.append(corners) else: # regular rooms if len(corners) >= 3 and Polygon(corners).area >= 100: room_polys.append(corners) room_types.append(pred_room_label_per_scene[j]) # window / door elif len(corners) == 2: window_doors.append(corners) window_doors_types.append(pred_room_label_per_scene[j]) if not semantic_rich: pred_room_label_per_scene = len(room_polys) * [-1] else: # poly2seq mode: sequence-based processing all_room_polys = [] tmp = [] all_length_list = [0] for j in range(len(pred_corners_per_scene)): if isinstance(pred_corners_per_scene[j], int): if pred_corners_per_scene[j] == 2 and tmp: # sep all_room_polys.append(tmp) all_length_list.append(len(tmp) + 1 + add_cls_token) tmp = [] continue tmp.append(pred_corners_per_scene[j]) if len(tmp): all_room_polys.append(tmp) all_length_list.append(len(tmp) + 1 + add_cls_token) start_poly_indices = np.cumsum(all_length_list) final_pred_classes = [] for j, poly in enumerate(all_room_polys): if len(poly) < 2: continue corners = np.array(poly, dtype=np.float32) * (image_size - 1) corners = np.around(corners).astype(np.int32) if not semantic_rich: # only regular rooms if len(corners) >= 4 and Polygon(corners).area >= 100: room_polys.append(corners) else: if per_token_sem_loss: pred_classes, counts = np.unique( pred_room_label_per_scene[start_poly_indices[j] : start_poly_indices[j + 1]][:-1], return_counts=True, ) pred_class = pred_classes[np.argmax(counts)] pred_logit = pred_room_logit_per_scene[start_poly_indices[j] : start_poly_indices[j + 1]][:-1] else: pred_class = pred_room_label_per_scene[ start_poly_indices[j + 1] - 1 ] # get last cls token in the seq final_pred_classes.append(pred_class) if wd_as_line: # regular rooms if len(corners) >= 3 and Polygon(corners).area >= 100: room_polys.append(corners) room_types.append(pred_class) # window / door elif len(corners) == 2: window_doors.append(corners) if ( door_window_index is not None and pred_class not in door_window_index and dataset_name != "r2g" ): wd_prob = np_softmax(pred_logit[:, door_window_index].sum(0)) pred_class = door_window_index[wd_prob.argmax()] window_doors_types.append(pred_class) else: # regular rooms if door_window_index is not None and pred_class not in door_window_index: room_polys.append(corners) room_types.append(pred_class) else: window_doors.append(corners) window_doors_types.append(pred_class) if not semantic_rich: pred_room_label_per_scene = len(all_room_polys) * [-1] else: pred_room_label_per_scene = final_pred_classes return { "room_polys": room_polys, "room_types": room_types if semantic_rich else None, "window_doors": window_doors if semantic_rich else None, "window_doors_types": window_doors_types if semantic_rich else None, "pred_room_label_per_scene": pred_room_label_per_scene, } @torch.no_grad() def evaluate_floor( model, dataset_name, data_loader, device, output_dir, plot_pred=True, plot_density=True, plot_gt=True, semantic_rich=False, save_pred=False, poly2seq: bool = False, add_cls_token=False, per_token_sem_loss=False, iou_thres=0.5, wd_as_line=True, ): """ Evaluate the model on a given dataset at testing. This function supports two evaluation modes: - Non-poly2seq mode (poly2seq=False): RoomFormer evaluation with fg_mask filtering - Poly2seq mode (poly2seq=True): Raster2Seq evaluation with sequence-based predictions This function processes the dataset in batches, performs inference using the provided model, and computes quantitative metrics such as precision, recall, and F1-score for rooms, corners, and angles. It supports semantic-rich evaluation including room types and window/door detection for specific datasets. Optionally, it can plot predicted floorplans, density maps, and save predictions as JSON files. Args: model: The trained model to evaluate (e.g., a PyTorch model). dataset_name (str): Name of the dataset (e.g., "stru3d", "cubicasa", "r2g", "waffle"). data_loader: PyTorch DataLoader providing batched inputs and extras. device: Device to run the model on (e.g., torch.device('cuda') or 'cpu'). output_dir (str): Directory path to save output files (plots, JSONs, results). plot_pred (bool, optional): If True, plot and save predicted floorplans. Default is True. plot_density (bool, optional): If True, plot and save density maps. Default is True. plot_gt (bool, optional): If True, save ground truth images. Default is True. semantic_rich (bool, optional): If True, perform semantic evaluation including room types and window/door lines. Default is False. save_pred (bool, optional): If True, save predicted polygons and metrics as JSON files. Default is False. poly2seq (bool, optional): If True, uses sequence-based prediction mode with forward_inference. If False, uses mask-based filtering with fg_mask. Defaults to False. add_cls_token (bool, optional): If True, accounts for class tokens in sequence processing. Only used when poly2seq=True. Defaults to False. per_token_sem_loss (bool, optional): If True, computes semantic loss per token using voting. Only used when poly2seq=True. Defaults to False. iou_thres (float, optional): IoU threshold for matching predictions to ground truth. Default is 0.5. wd_as_line (bool, optional): If True, treats windows/doors as lines based on corner count. If False, uses semantic class. Only used when poly2seq=True. Defaults to True. Returns: None. The function prints the aggregated quantitative results to the console and saves them to 'results.txt' in output_dir. """ model.eval() if dataset_name == "stru3d": door_window_index = [16, 17] elif dataset_name == "cubicasa": door_window_index = [10, 9] elif dataset_name == "waffle": door_window_index = [] # [1, 2] else: door_window_index = [] metric_category = ["room", "corner", "angles"] if semantic_rich: metric_category += ["room_sem", "window_door"] quant_result_dict = None scene_counter = 0 merge = False # Only used in poly2seq mode for merged plot output if not os.path.exists(output_dir): os.mkdir(output_dir) for batched_inputs, batched_extras in data_loader: samples = [x["image"].to(device) for x in batched_inputs] scene_ids = [x["image_id"] for x in batched_inputs] gt_instances = [x["instances"].to(device) for x in batched_inputs] image_size = samples[0].size(2) # Get predictions based on mode if poly2seq: outputs = model.forward_inference(samples) pred_corners = outputs["gen_out"] fg_mask = None else: outputs = model(samples) pred_logits = outputs["pred_logits"] pred_corners = outputs["pred_coords"] fg_mask = torch.sigmoid(pred_logits) > 0.5 # select valid corners bs = outputs["pred_logits"].shape[0] if "pred_room_logits" in outputs: pred_room_logits = outputs["pred_room_logits"] prob = torch.nn.functional.softmax(pred_room_logits, -1) _, pred_room_label = prob[..., :-1].max(-1) # process per scene for i in range(bs): # Prepare ground truth data gt_polys, gt_polys_types = [], [] gt_window_doors = [] gt_window_doors_types = [] for gt_poly, gt_id in zip( gt_instances[i].gt_masks.polygons, gt_instances[i].gt_classes.detach().cpu().tolist() ): gt_poly = gt_poly[0].reshape(-1, 2).astype(np.int32) if gt_id in door_window_index: gt_window_doors.append(gt_poly) gt_window_doors_types.append(gt_id) else: gt_polys.append(gt_poly) gt_polys_types.append(gt_id) # Create evaluator based on dataset if dataset_name == "stru3d": if int(scene_ids[i]) in wrong_s3d_annotations_list: continue curr_opts = copy.deepcopy(opts) curr_opts.scene_id = "scene_0" + str(scene_ids[i]) curr_data_rw = S3DRW(curr_opts, mode="test") evaluator = Evaluator(curr_data_rw, curr_opts, disable_overlapping_filter=poly2seq) elif dataset_name in ["cubicasa", "waffle", "r2g"]: evaluator = Evaluator_RPlan( disable_overlapping_filter=poly2seq, iou_thres=iou_thres, wd_as_line=wd_as_line ) print("Running Evaluation for scene %s" % scene_ids[i]) scene_outputs = _process_predictions( pred_corners, i, semantic_rich, poly2seq, fg_mask, image_size, pred_room_label if semantic_rich else None, pred_room_logits if semantic_rich else None, add_cls_token, per_token_sem_loss, wd_as_line, door_window_index, dataset_name, ) room_polys = scene_outputs["room_polys"] room_types = scene_outputs["room_types"] window_doors = scene_outputs["window_doors"] window_doors_types = scene_outputs["window_doors_types"] pred_room_label_per_scene = scene_outputs["pred_room_label_per_scene"] if dataset_name == "stru3d": if not semantic_rich: quant_result_dict_scene = evaluator.evaluate_scene(room_polys=room_polys) else: quant_result_dict_scene = evaluator.evaluate_scene( room_polys=room_polys, room_types=room_types, window_door_lines=window_doors, window_door_lines_types=window_doors_types, ) elif dataset_name in ["cubicasa", "waffle", "r2g"]: if not semantic_rich: quant_result_dict_scene = evaluator.evaluate_scene( room_polys=room_polys, gt_polys=gt_polys, room_types=None, gt_polys_types=gt_polys_types, img_size=(image_size, image_size), ) else: quant_result_dict_scene = evaluator.evaluate_scene( room_polys=room_polys, gt_polys=gt_polys, room_types=room_types, gt_polys_types=gt_polys_types, window_door_lines=window_doors, gt_window_doors_list=gt_window_doors, window_door_lines_types=window_doors_types, gt_window_doors_type_list=gt_window_doors_types, img_size=(image_size, image_size), ) if quant_result_dict is None: quant_result_dict = quant_result_dict_scene else: for k in quant_result_dict.keys(): quant_result_dict[k] += quant_result_dict_scene[k] scene_counter += 1 # plot regular room floorplan gt_room_polys = [np.array(poly) for poly in gt_polys] room_polys = [np.array(poly) for poly in room_polys] if "gt_polys_sorted_indcs" in quant_result_dict_scene: gt_polys_sorted_indcs = quant_result_dict_scene["gt_polys_sorted_indcs"] del quant_result_dict_scene["gt_polys_sorted_indcs"] gt_room_polys = [gt_room_polys[ind] for ind in gt_polys_sorted_indcs] if "pred2gt_indices" in quant_result_dict_scene: pred2gt_indices = quant_result_dict_scene["pred2gt_indices"] del quant_result_dict_scene["pred2gt_indices"] room_polys, gt_room_polys, pred_mask, gt_mask = sort_polygons_by_matching( pred2gt_indices, room_polys, gt_room_polys ) else: pred_mask, gt_mask = None, None prec, rec = quant_result_dict_scene["room_prec"], quant_result_dict_scene["room_rec"] f1 = 2 * prec * rec / (prec + rec + 1e-5) missing_rate = quant_result_dict_scene["room_missing_ratio"] plot_statistics = { "f1": f1, "prec": prec, "rec": rec, "missing_rate": missing_rate, "num_preds": len(room_polys), "num_gt": len(gt_polys), "num_matched_preds": sum([x != -1 for x in pred2gt_indices]), } if plot_pred: gt_floorplan_map = plot_floorplan_with_regions( gt_room_polys, matching_labels=gt_mask, base_scale=image_size, scale=1024 ) floorplan_map = plot_floorplan_with_regions( room_polys, matching_labels=pred_mask, base_scale=image_size, scale=1024 ) if not merge: cv2.imwrite(os.path.join(output_dir, "{}_pred_floorplan.png".format(scene_ids[i])), floorplan_map) cv2.imwrite(os.path.join(output_dir, "{}_gt_floorplan.png".format(scene_ids[i])), gt_floorplan_map) else: concatenated_map = concat_floorplan_maps(gt_floorplan_map, floorplan_map, plot_statistics) cv2.imwrite( os.path.join(output_dir, "{}_pred_floorplan.png".format(scene_ids[i])), concatenated_map ) if semantic_rich: _, ID2CLASS_LABEL = get_dataset_class_labels(dataset_name) floorplan_map = plot_semantic_rich_floorplan_opencv( zip(room_polys + window_doors, room_types + window_doors_types), os.path.join(output_dir, "{}_pred_floorplan_sem.png".format(scene_ids[i])), door_window_index=door_window_index, semantics_label_mapping=ID2CLASS_LABEL, img_w=image_size, img_h=image_size, scale=1, plot_text=False, ) if save_pred: # Save room_polys as JSON json_path = os.path.join(output_dir, "jsons", "{}.json".format(str(scene_ids[i]).zfill(5))) os.makedirs(os.path.dirname(json_path), exist_ok=True) polys_list = [poly.astype(float).tolist() for poly in room_polys] if semantic_rich: polys_list += [window_door.astype(float).tolist() for window_door in window_doors] types_list = room_types + window_doors_types else: types_list = [-1] * len(polys_list) output_json = [ { "image_id": str(scene_ids[i]).zfill(5), "segmentation": polys_list[instance_id], "category_id": int(types_list[instance_id]), "id": instance_id, } for instance_id in range(len(polys_list)) ] with open(json_path, "w") as json_file: json.dump(output_json, json_file) json_result_path = os.path.join( output_dir, "result_jsons", "{}.json".format(str(scene_ids[i]).zfill(5)) ) new_quant_result_dict_scene = compute_f1(copy.deepcopy(quant_result_dict_scene), metric_category) os.makedirs(os.path.dirname(json_result_path), exist_ok=True) with open(json_result_path, "w") as json_file: json.dump(new_quant_result_dict_scene, json_file) if plot_gt: gt_image = np.transpose(samples[i].cpu().numpy(), (1, 2, 0)) gt_image = (gt_image * 255).astype(np.uint8) cv2.imwrite(os.path.join(output_dir, "{}_gt.png".format(scene_ids[i])), gt_image) if plot_density: pred_room_map = plot_density_map( samples[i], image_size, room_polys, pred_room_label_per_scene, plot_text=False ) gt_room_map = plot_density_map(samples[i], image_size, gt_polys, gt_polys_types, plot_text=False) if merge: concatenated_map = concat_floorplan_maps(gt_room_map, pred_room_map, plot_statistics) cv2.imwrite( os.path.join(output_dir, "{}_pred_room_map.png".format(scene_ids[i])), concatenated_map ) else: cv2.imwrite(os.path.join(output_dir, "{}_pred_room_map.png".format(scene_ids[i])), pred_room_map) cv2.imwrite(os.path.join(output_dir, "{}_gt_room_map.png".format(scene_ids[i])), gt_room_map) for k in quant_result_dict.keys(): quant_result_dict[k] /= float(scene_counter) quant_result_dict = compute_f1(quant_result_dict, metric_category) print("*************************************************") print(quant_result_dict) print("*************************************************") with open(os.path.join(output_dir, "results.txt"), "w") as file: file.write(json.dumps(quant_result_dict)) def generate( model, samples, semantic_rich=False, poly2seq: bool = False, use_cache=True, per_token_sem_loss=False, drop_wd=False, ): """ Generate room polygons and labels from model predictions. This function supports two generation modes: - Non-poly2seq mode (poly2seq=False): RoomFormer generation with fg_mask filtering - Poly2seq mode (poly2seq=True): Raster2Seq generation with sequence-based predictions Args: model: The trained model to use for inference. samples: Input image samples (list of tensors). semantic_rich (bool, optional): If True, predict room types and window/door elements. Default is False. poly2seq (bool, optional): If True, uses sequence-based prediction mode with forward_inference. If False, uses mask-based filtering with fg_mask. Defaults to False. use_cache (bool, optional): If True, use caching in forward_inference. Only used when poly2seq=True. Default is True. per_token_sem_loss (bool, optional): If True, computes semantic class using voting across tokens. Only used when poly2seq=True. Default is False. drop_wd (bool, optional): If True, exclude window/door elements from output. Default is False. Returns: dict: Dictionary containing: - 'room': List of room polygons (and optionally window/door lines) per scene - 'labels': List of class labels per scene """ model.eval() image_size = samples[0].size(2) # Get predictions based on mode if poly2seq: outputs = model.forward_inference(samples, use_cache) pred_corners = outputs["gen_out"] fg_mask = None else: outputs = model(samples) pred_corners = outputs["pred_coords"] pred_logits = outputs["pred_logits"] fg_mask = torch.sigmoid(pred_logits) > 0.5 # select valid corners bs = outputs["pred_logits"].shape[0] if "pred_room_logits" in outputs: pred_room_logits = outputs["pred_room_logits"] prob = torch.nn.functional.softmax(pred_room_logits, -1) _, pred_room_label = prob[..., :-1].max(-1) outputs = [] output_classes = [] # process per scene for i in range(bs): room_polys = [] if semantic_rich: room_types = [] window_doors = [] window_doors_types = [] else: window_doors = None room_types = None scene_outputs = _process_predictions( pred_corners, i, semantic_rich, poly2seq, fg_mask, image_size, pred_room_label if semantic_rich else None, pred_room_logits if semantic_rich else None, False, per_token_sem_loss, wd_as_line=True, door_window_index=None, dataset_name=None, ) room_polys = scene_outputs["room_polys"] room_types = scene_outputs["room_types"] window_doors = scene_outputs["window_doors"] window_doors_types = scene_outputs["window_doors_types"] if not drop_wd and window_doors: outputs.append(room_polys + window_doors) output_classes.append(room_types + window_doors_types) else: outputs.append(room_polys) output_classes.append(room_types) out_dict = {"room": outputs, "labels": output_classes} return out_dict