File size: 1,651 Bytes
61b3c2b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from pathlib import Path
from typing import Dict
import shutil
from PIL import Image
import glob
import tempfile
from sorghum_pipeline.pipeline import SorghumPipeline
def run_pipeline_on_image(input_image_path: str, work_dir: str, save_artifacts: bool = True) -> Dict[str, str]:
"""
Run sorghum pipeline on a single image (no instance segmentation).
Returns dict[label -> image_path] for gallery display.
"""
work = Path(work_dir)
work.mkdir(parents=True, exist_ok=True)
# Copy input to work dir
input_copy = work / Path(input_image_path).name
shutil.copy(input_image_path, input_copy)
# Initialize pipeline with config
# adjust this if you have a YAML config file (e.g., "configs/demo.yaml")
pipeline = SorghumPipeline(
config_path=str(Path("sorghum_pipeline/config.py")),
enable_occlusion_handling=False,
enable_instance_integration=False
)
# Run the pipeline (single image, no frames, no SAM2Long)
results = pipeline.run(
load_all_frames=False,
segmentation_only=False,
run_instance_segmentation=False,
features_frame_only=None
)
# Collect outputs
outputs: Dict[str, str] = {}
# Save original for reference
original = work / "original.png"
Image.open(input_copy).convert("RGB").save(original)
outputs["Original"] = str(original)
# Gather all PNG files created by OutputManager
for f in glob.glob(str(work / "**/*.png"), recursive=True):
name = Path(f).stem
if name.lower() not in outputs: # avoid duplicate "Original"
outputs[name] = f
return outputs
|