update sorghum for multiple plants

app.py

@@ -23,11 +23,14 @@ PRESET_IMAGES = {
 
 def process(file_path, preset_choice):
     """Process image and yield results progressively for immediate display."""
-    # If a preset is chosen, override the uploaded file path
+    # Determine dataset type (single-plant mode for Corn, multi-plant for others)
+    single_plant_mode = False
     if preset_choice:
         chosen = PRESET_IMAGES.get(preset_choice)
         if chosen:
             file_path = chosen
+            # Corn uses single-plant mode
+            single_plant_mode = (preset_choice == "Corn")
 
     if not file_path:
         # Return 10 outputs (removed YOLO tips)
@@ -62,7 +65,7 @@ def process(file_path, preset_choice):
                 return None
 
         # Run the pipeline progressively (generator)
-        for outputs in run_pipeline_on_image(str(img_path), tmpdir, save_artifacts=True):
+        for outputs in run_pipeline_on_image(str(img_path), tmpdir, save_artifacts=True, single_plant_mode=single_plant_mode):
             # Load all available outputs progressively
             composite = load_pil(outputs.get('Composite'))
             overlay = load_pil(outputs.get('Overlay'))

sorghum_pipeline/features/morphology.py

@@ -22,12 +22,13 @@ class MorphologyExtractor:
     """Morphology extraction: size analysis image + simple traits."""
     
     def __init__(self, pixel_to_cm: float = 0.1099609375, prune_sizes: List[int] = None,
-                 min_component_area_for_size: int = 3000):
+                 min_component_area_for_size: int = 3000, single_plant_mode: bool = False):
         """Initialize."""
         self.pixel_to_cm = pixel_to_cm
         self.prune_sizes = prune_sizes or [200, 100, 50, 30, 10]
         # Used only for the Morphology Size visualization (not for height)
         self.min_component_area_for_size = int(min_component_area_for_size)
+        self.single_plant_mode = single_plant_mode
         
         if PLANT_CV_AVAILABLE:
             pcv.params.debug = None
@@ -72,19 +73,22 @@ class MorphologyExtractor:
                     plant_heights[f'plant_{plant_idx}'] = height_cm
                     plant_areas[f'plant_{plant_idx}'] = area
             
-            # Keep only the largest plant (main plant) for single-plant datasets
-            if len(plant_heights) > 1:
+            # Keep only the largest plant if in single-plant mode (e.g., corn)
+            if self.single_plant_mode and len(plant_heights) > 1:
                 # Find the largest plant by area
                 largest_plant = max(plant_areas.items(), key=lambda x: x[1])[0]
                 plant_heights = {largest_plant: plant_heights[largest_plant]}
             
             # Store individual plant heights
             features['traits']['plant_heights'] = plant_heights
-            features['traits']['num_plants'] = 1 if len(plant_heights) > 0 else 0
+            features['traits']['num_plants'] = len(plant_heights) if not self.single_plant_mode else (1 if len(plant_heights) > 0 else 0)
             
-            # Store single plant height
+            # For backward compatibility, store single height
             if len(plant_heights) == 1:
                 features['traits']['plant_height_cm'] = list(plant_heights.values())[0]
+            elif len(plant_heights) > 1:
+                # Store max height as overall height for multi-plant
+                features['traits']['plant_height_cm'] = max(plant_heights.values())
             else:
                 features['traits']['plant_height_cm'] = 0.0
 
@@ -165,43 +169,54 @@ class MorphologyExtractor:
         return arr
     
     def _simple_size_visual(self, rgb: np.ndarray, mask: np.ndarray) -> np.ndarray:
-        """Draw contours and bbox for the largest plant on RGB image."""
+        """Draw contours and bbox for plants on RGB image."""
         vis = rgb.copy()
         
         # Find connected components to identify individual plants
         num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
         
-        # Find the largest plant (skip background 0)
-        largest_idx = -1
-        largest_area = 0
-        for plant_idx in range(1, num_labels):
-            area = stats[plant_idx, cv2.CC_STAT_AREA]
-            if area > largest_area and area >= 100:  # Filter noise
-                largest_area = area
-                largest_idx = plant_idx
-        
-        # Draw only the largest plant
-        if largest_idx > 0:
-            # Get mask for largest plant
-            plant_mask = ((labels == largest_idx).astype(np.uint8) * 255)
-            
-            # Find contours
-            contours, _ = cv2.findContours(plant_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            
-            # Use blue color for main plant
-            color = (255, 0, 0)
+        if self.single_plant_mode:
+            # Single-plant mode: draw only the largest plant
+            largest_idx = -1
+            largest_area = 0
+            for plant_idx in range(1, num_labels):
+                area = stats[plant_idx, cv2.CC_STAT_AREA]
+                if area > largest_area and area >= 100:  # Filter noise
+                    largest_area = area
+                    largest_idx = plant_idx
             
-            # Draw contours
-            cv2.drawContours(vis, contours, -1, color, 2)
+            if largest_idx > 0:
+                plant_mask = ((labels == largest_idx).astype(np.uint8) * 255)
+                contours, _ = cv2.findContours(plant_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                color = (255, 0, 0)
+                cv2.drawContours(vis, contours, -1, color, 2)
+                if contours:
+                    x, y, w, h = cv2.boundingRect(contours[0])
+                    cv2.rectangle(vis, (x, y), (x + w, y + h), (0, 255, 0), 2)
+                    cv2.putText(vis, "Plant 1", (x, y - 5), 
+                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
+        else:
+            # Multi-plant mode: draw all plants with different colors
+            colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255)]
+            plant_count = 0
             
-            # Draw bounding box
-            if contours:
-                x, y, w, h = cv2.boundingRect(contours[0])
-                cv2.rectangle(vis, (x, y), (x + w, y + h), (0, 255, 0), 2)
+            for plant_idx in range(1, num_labels):
+                area = stats[plant_idx, cv2.CC_STAT_AREA]
+                if area < 100:  # Filter noise
+                    continue
+                
+                plant_mask = ((labels == plant_idx).astype(np.uint8) * 255)
+                contours, _ = cv2.findContours(plant_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                color = colors[plant_count % len(colors)]
+                cv2.drawContours(vis, contours, -1, color, 2)
+                
+                if contours:
+                    x, y, w, h = cv2.boundingRect(contours[0])
+                    cv2.rectangle(vis, (x, y), (x + w, y + h), color, 2)
+                    cv2.putText(vis, f"P{plant_idx}", (x, y - 5), 
+                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
                 
-                # Add "Plant 1" label
-                cv2.putText(vis, "Plant 1", (x, y - 5), 
-                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
+                plant_count += 1
         
         return vis
     

sorghum_pipeline/pipeline.py

@@ -20,18 +20,19 @@ logger = logging.getLogger(__name__)
 class SorghumPipeline:
     """Minimal pipeline for single-image processing."""
     
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, single_plant_mode: bool = False):
         """Initialize pipeline."""
         logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
         self.config = config
         self.config.validate()
+        self.single_plant_mode = single_plant_mode
         
         # Initialize components with defaults
         self.preprocessor = ImagePreprocessor()
         self.mask_handler = MaskHandler()
         self.texture_extractor = TextureExtractor()
         self.vegetation_extractor = VegetationIndexExtractor()
-        self.morphology_extractor = MorphologyExtractor()
+        self.morphology_extractor = MorphologyExtractor(single_plant_mode=single_plant_mode)
         self.segmentation_manager = SegmentationManager(
             model_name="briaai/RMBG-2.0",
             device=self.config.get_device(),
@@ -41,7 +42,7 @@ class SorghumPipeline:
             output_folder=self.config.paths.output_folder,
             settings=self.config.output
         )
-        logger.info("Pipeline initialized")
+        logger.info(f"Pipeline initialized (single_plant_mode={single_plant_mode})")
 
     def run(self, single_image_path: str) -> Dict[str, Any]:
         """Run pipeline on single image."""
@@ -105,7 +106,7 @@ class SorghumPipeline:
 
         return {"plants": plants, "timing": elapsed}
 
-    def run_with_progress(self, single_image_path: str, progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None) -> Generator[Dict[str, Any], None, None]:
+    def run_with_progress(self, single_image_path: str, progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, single_plant_mode: bool = False) -> Generator[Dict[str, Any], None, None]:
         """Run pipeline on single image, yielding intermediate results progressively."""
         logger.info("Processing single image with progress...")
 

wrapper.py

@@ -10,7 +10,8 @@ from sorghum_pipeline.config import Config, Paths
 
 
 def run_pipeline_on_image(input_image_path: str, work_dir: str, save_artifacts: bool = True, 
-                          progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None) -> Generator[Dict[str, str], None, None]:
+                          progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None,
+                          single_plant_mode: bool = False) -> Generator[Dict[str, str], None, None]:
     """
     Run sorghum pipeline on a single image (no instance segmentation).
     Yields dict[label -> image_path] progressively for gallery display.
@@ -42,10 +43,10 @@ def run_pipeline_on_image(input_image_path: str, work_dir: str, save_artifacts:
         output_folder=str(work),
         boundingbox_dir=str(work)
     )
-    pipeline = SorghumPipeline(config=cfg)
+    pipeline = SorghumPipeline(config=cfg, single_plant_mode=single_plant_mode)
 
     # Run the pipeline with progress callback (generator)
-    for stage_result in pipeline.run_with_progress(single_image_path=str(input_path), progress_callback=progress_callback):
+    for stage_result in pipeline.run_with_progress(single_image_path=str(input_path), progress_callback=progress_callback, single_plant_mode=single_plant_mode):
         # Yield intermediate outputs as they become available
         outputs = _collect_outputs(work, stage_result.get('plants', {}))
         yield outputs
@@ -114,11 +115,20 @@ def _collect_outputs(work: Path, plants: Dict[str, Any]) -> Dict[str, str]:
             plant_heights = traits.get('plant_heights', {})
             num_plants = traits.get('num_plants', 0)
             
-            # Always show as single plant (largest component)
-            if num_plants > 0 and isinstance(plant_heights, dict) and len(plant_heights) >= 1:
-                height_cm = list(plant_heights.values())[0]
-                stats_lines.append(f"Number of plants: 1")
-                stats_lines.append(f"Plant height: {height_cm:.2f} cm")
+            # Display plant info based on mode
+            if num_plants > 0 and isinstance(plant_heights, dict):
+                if num_plants == 1 or len(plant_heights) == 1:
+                    # Single plant display
+                    height_cm = list(plant_heights.values())[0]
+                    stats_lines.append(f"Number of plants: 1")
+                    stats_lines.append(f"Plant height: {height_cm:.2f} cm")
+                else:
+                    # Multiple plants display
+                    stats_lines.append(f"Number of plants: {num_plants}")
+                    sorted_plants = sorted(plant_heights.items(), key=lambda x: int(x[0].split('_')[1]))
+                    for plant_name, height_cm in sorted_plants:
+                        plant_num = plant_name.split('_')[1]
+                        stats_lines.append(f"  Plant {plant_num}: {height_cm:.2f} cm")
             else:
                 # Fallback to old single height field
                 height_cm = traits.get('plant_height_cm')