Upload infer/infer.py with huggingface_hub

infer/infer.py

@@ -17,7 +17,12 @@ if PROJECT_ROOT not in sys.path:
 logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.ERROR)
 os.environ["CUDA_VISIBLE_DEVICES"] = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
 
-from infer.common_infer import initialize_pipeline, quantize_box_16, scale_box_xyxy
+from infer.common_infer import (
+    get_layer_boxes,
+    initialize_pipeline,
+    quantize_box_16,
+    scale_box_xyxy,
+)
 from tools.tools import load_config, seed_everything
 
 
@@ -77,15 +82,27 @@ def build_run_save_dir(config: dict):
 
 
 def resolve_image_path(sample: dict, data_dir: str, image_dir: str = None) -> str:
-    """Resolve the input image path, preferring local files_real_test images."""
+    """Resolve the input image path for Stage 1 or Prism-style metadata."""
+    sample_dir = sample.get("sample_dir", "")
     sample_name = sample.get("sample_or_stem", "")
     image_path = sample.get("image", "")
+    blend_path = sample.get("blend_path", "")
 
     if image_dir is None and data_dir:
         image_dir = os.path.join(data_dir, "layers_real_test_1024")
 
     candidates = []
 
+    if sample_dir:
+        if data_dir and not os.path.isabs(sample_dir):
+            candidates.append(os.path.join(data_dir, sample_dir, "whole_image.png"))
+        candidates.append(os.path.join(sample_dir, "whole_image.png"))
+
+    if blend_path:
+        candidates.append(blend_path)
+        if data_dir and not os.path.isabs(blend_path):
+            candidates.append(os.path.join(data_dir, blend_path))
+
     if image_dir:
         if sample_name:
             candidates.extend(
@@ -139,7 +156,10 @@ def quantize_box_16_safe(box: tuple, target_size: int) -> tuple:
 
 
 def get_real_boxes(sample: dict, source_size: int, target_size: int) -> list:
-    """Scale and quantize real-test boxes from JSON metadata."""
+    """Scale and quantize boxes from Stage 1 or Prism-style JSON metadata."""
+    if sample.get("layers"):
+        return get_layer_boxes(sample.get("layers", []), source_size, target_size)
+
     boxes = []
     for box in sample.get("bboxes", []):
         if not isinstance(box, (list, tuple)) or len(box) != 4:
@@ -228,7 +248,11 @@ def inference_real(config):
 
     for local_idx, sample in enumerate(samples):
         idx_zero_based = start_idx - 1 + local_idx
-        sample_name = sample.get("sample_or_stem", f"real_{idx_zero_based:06d}")
+        sample_name = (
+            sample.get("sample_or_stem")
+            or sample.get("sample_dir")
+            or f"sample_{idx_zero_based:06d}"
+        )
         print(
             f"Processing [{local_idx + 1}/{len(samples)}] idx={idx_zero_based} ({sample_name})...",
             flush=True,
@@ -316,9 +340,9 @@ def inference_real(config):
             "source_image_path": format_source_image_path(image_path, config),
             "target_size": target_size,
             "source_size": source_size,
-            "raw_num_layers": sample.get("num_layers"),
+            "raw_num_layers": sample.get("num_layers", sample.get("layer_count")),
             "num_layers": len(all_boxes),
-            "raw_boxes": sample.get("bboxes", []),
+            "raw_boxes": sample.get("bboxes", [layer.get("box") for layer in sample.get("layers", [])]),
             "boxes": all_boxes,
             "caption": caption,
             "run_name": run_name,