Eduardo Ottone
ee92581211
- Implements YOLOv8-based object detection with batch output - Supports multiple YOLOv8 model sizes (nano to extra-large) - Configurable confidence threshold, IoU threshold, and padding - Outputs separate images for each detected object - Includes installation scripts for Windows, Linux, and Mac - Complete documentation and example workflow - MIT license for open source distribution
570 lines
22 KiB
Python
570 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Practical segmentation for useful items only
|
|
Segments: hair, head/body, individual clothes, accessories, and scene objects
|
|
NO tiny body parts like ears, eyes, etc.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
from ultralytics import YOLO
|
|
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
|
|
from typing import List, Dict, Any, Tuple
|
|
import supervision as sv
|
|
from scipy import ndimage
|
|
|
|
class PracticalSegmenter:
|
|
def __init__(self, sam_model='vit_b', yolo_model='yolov8m.pt', device='auto'):
|
|
"""
|
|
Initialize practical segmentation
|
|
"""
|
|
# Auto-detect device
|
|
if device == 'auto':
|
|
if torch.cuda.is_available():
|
|
self.device = 'cuda'
|
|
self.yolo_device = 'cuda'
|
|
print("Using CUDA (GPU)")
|
|
elif torch.backends.mps.is_available():
|
|
self.device = 'cpu' # SAM on CPU
|
|
self.yolo_device = 'mps' # YOLO on MPS
|
|
print("Using CPU for SAM, MPS for YOLO")
|
|
else:
|
|
self.device = 'cpu'
|
|
self.yolo_device = 'cpu'
|
|
print("Using CPU")
|
|
else:
|
|
self.device = device
|
|
self.yolo_device = device
|
|
|
|
# Load YOLO for object detection
|
|
print("Loading YOLO for object detection...")
|
|
self.yolo = YOLO(yolo_model)
|
|
self.yolo.to(self.yolo_device)
|
|
|
|
# Load SAM
|
|
sam_path = Path(f'sam_{sam_model}.pth')
|
|
if sam_path.exists():
|
|
print(f"Loading SAM {sam_model} model...")
|
|
self.sam = sam_model_registry[sam_model](checkpoint=str(sam_path))
|
|
self.sam.to(device=self.device)
|
|
|
|
# Configure for practical segmentation
|
|
self.mask_generator = SamAutomaticMaskGenerator(
|
|
self.sam,
|
|
points_per_side=32, # Balanced for speed and quality
|
|
pred_iou_thresh=0.86, # Good threshold
|
|
stability_score_thresh=0.92,
|
|
crop_n_layers=1, # Single crop layer for speed
|
|
crop_n_points_downscale_factor=2,
|
|
min_mask_region_area=100, # Small minimum to catch details
|
|
)
|
|
else:
|
|
print("SAM model not found. Run segment_everything.py first to download.")
|
|
sys.exit(1)
|
|
|
|
# Categories we actually want
|
|
self.wanted_categories = {
|
|
'person_parts': ['hair', 'head', 'face', 'body', 'person'],
|
|
'clothing': ['shirt', 'jacket', 'coat', 'pants', 'dress', 'skirt', 'shoes', 'hat'],
|
|
'accessories': ['bag', 'backpack', 'glasses', 'watch', 'jewelry', 'belt'],
|
|
'scene_objects': ['furniture', 'vase', 'plant', 'book', 'device', 'decoration']
|
|
}
|
|
|
|
def segment_practical(self, image_path: str, merge_body_parts: bool = True) -> Dict[str, Any]:
|
|
"""
|
|
Segment image into practical, useful segments
|
|
"""
|
|
# Read image
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
raise ValueError(f"Could not read image: {image_path}")
|
|
|
|
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
height, width = image.shape[:2]
|
|
|
|
print(f"\nProcessing {Path(image_path).name} for practical segmentation...")
|
|
|
|
results = {
|
|
'original_image': image_rgb,
|
|
'image_path': image_path,
|
|
'image_size': (width, height),
|
|
'segments': {
|
|
'hair': [],
|
|
'head_face': [],
|
|
'body': [],
|
|
'clothing': [],
|
|
'accessories': [],
|
|
'scene_objects': []
|
|
}
|
|
}
|
|
|
|
# Step 1: Detect objects with YOLO
|
|
print("Detecting main objects...")
|
|
yolo_results = self.yolo(image_rgb, conf=0.25)
|
|
|
|
person_boxes = []
|
|
other_objects = []
|
|
|
|
if len(yolo_results) > 0 and yolo_results[0].boxes is not None:
|
|
for i, box in enumerate(yolo_results[0].boxes):
|
|
class_id = int(box.cls.cpu().numpy()[0])
|
|
class_name = self.yolo.names[class_id]
|
|
x1, y1, x2, y2 = box.xyxy.cpu().numpy()[0].astype(int)
|
|
conf = float(box.conf.cpu().numpy()[0])
|
|
|
|
if class_name == 'person':
|
|
person_boxes.append({
|
|
'bbox': (x1, y1, x2, y2),
|
|
'confidence': conf
|
|
})
|
|
else:
|
|
# Other scene objects
|
|
other_objects.append({
|
|
'type': class_name,
|
|
'bbox': (x1, y1, x2, y2),
|
|
'confidence': conf
|
|
})
|
|
|
|
print(f"Found {len(person_boxes)} person(s) and {len(other_objects)} other objects")
|
|
|
|
# Step 2: Generate SAM masks
|
|
print("Generating detailed segments...")
|
|
all_masks = self.mask_generator.generate(image_rgb)
|
|
|
|
# Sort by area (largest first)
|
|
all_masks = sorted(all_masks, key=lambda x: x['area'], reverse=True)
|
|
|
|
print(f"Generated {len(all_masks)} initial segments")
|
|
|
|
# Step 3: Classify and merge segments
|
|
processed_masks = []
|
|
used_pixels = np.zeros((height, width), dtype=bool)
|
|
|
|
print(f"Processing {len(all_masks)} segments...")
|
|
kept_count = 0
|
|
skipped_small = 0
|
|
|
|
for mask_data in all_masks:
|
|
mask = mask_data['segmentation']
|
|
bbox = mask_data['bbox'] # x, y, w, h
|
|
x, y, w, h = [int(v) for v in bbox]
|
|
area = mask_data['area']
|
|
|
|
# Don't skip based on overlap - we want all segments!
|
|
|
|
# Classify the segment
|
|
segment_type = self._classify_segment(
|
|
mask, bbox, person_boxes, image_rgb, area
|
|
)
|
|
|
|
# Skip only very tiny segments unless they're important
|
|
if area < 100 and segment_type not in ['hair', 'head_face', 'accessory']:
|
|
skipped_small += 1
|
|
continue
|
|
|
|
kept_count += 1
|
|
|
|
# Process based on type
|
|
if segment_type in ['hair', 'head_face'] and merge_body_parts:
|
|
# Try to merge with nearby similar segments
|
|
mask = self._merge_nearby_segments(mask, all_masks, segment_type, bbox)
|
|
|
|
# Create segment info
|
|
isolated = image_rgb.copy()
|
|
isolated[mask == 0] = 0
|
|
|
|
# Crop to bounding box
|
|
y_min, y_max = np.where(mask.any(axis=1))[0][[0, -1]]
|
|
x_min, x_max = np.where(mask.any(axis=0))[0][[0, -1]]
|
|
|
|
isolated_cropped = isolated[y_min:y_max+1, x_min:x_max+1]
|
|
mask_cropped = mask[y_min:y_max+1, x_min:x_max+1]
|
|
|
|
segment_info = {
|
|
'type': segment_type,
|
|
'bbox': (x_min, y_min, x_max, y_max),
|
|
'area': np.sum(mask),
|
|
'mask': mask.astype(np.uint8),
|
|
'isolated_full': isolated,
|
|
'isolated_cropped': isolated_cropped,
|
|
'mask_cropped': mask_cropped.astype(np.uint8)
|
|
}
|
|
|
|
# Add to appropriate category
|
|
if segment_type == 'hair':
|
|
results['segments']['hair'].append(segment_info)
|
|
elif segment_type == 'head_face':
|
|
results['segments']['head_face'].append(segment_info)
|
|
elif segment_type == 'body':
|
|
results['segments']['body'].append(segment_info)
|
|
elif segment_type in ['shirt', 'jacket', 'pants', 'dress', 'skirt', 'shoes']:
|
|
segment_info['clothing_type'] = segment_type
|
|
results['segments']['clothing'].append(segment_info)
|
|
elif segment_type == 'accessory':
|
|
results['segments']['accessories'].append(segment_info)
|
|
elif segment_type == 'object':
|
|
results['segments']['scene_objects'].append(segment_info)
|
|
|
|
# Don't mark pixels as used since we want overlapping segments
|
|
# used_pixels |= mask
|
|
|
|
# Step 4: Add YOLO-detected scene objects
|
|
for obj in other_objects:
|
|
x1, y1, x2, y2 = obj['bbox']
|
|
|
|
# Create mask for object region
|
|
obj_mask = np.zeros((height, width), dtype=np.uint8)
|
|
obj_mask[y1:y2, x1:x2] = 1
|
|
|
|
# Extract object
|
|
isolated = image_rgb.copy()
|
|
isolated[obj_mask == 0] = 0
|
|
isolated_cropped = isolated[y1:y2, x1:x2]
|
|
|
|
segment_info = {
|
|
'type': obj['type'],
|
|
'bbox': obj['bbox'],
|
|
'area': (x2-x1) * (y2-y1),
|
|
'mask': obj_mask,
|
|
'isolated_full': isolated,
|
|
'isolated_cropped': isolated_cropped,
|
|
'mask_cropped': obj_mask[y1:y2, x1:x2],
|
|
'confidence': obj['confidence']
|
|
}
|
|
|
|
results['segments']['scene_objects'].append(segment_info)
|
|
|
|
# Print debug info
|
|
print(f"\nDebug: Kept {kept_count} segments, Skipped {skipped_small} small segments")
|
|
|
|
# Print summary
|
|
total_segments = sum(len(segs) for segs in results['segments'].values())
|
|
print(f"\nSegmentation complete! Found {total_segments} practical segments:")
|
|
for category, segments in results['segments'].items():
|
|
if segments:
|
|
print(f" {category}: {len(segments)} items")
|
|
|
|
return results
|
|
|
|
def _classify_segment(self, mask: np.ndarray, bbox: tuple, person_boxes: list,
|
|
image: np.ndarray, area: int) -> str:
|
|
"""
|
|
Classify a segment into practical categories
|
|
"""
|
|
x, y, w, h = bbox
|
|
height, width = image.shape[:2]
|
|
|
|
# Check if inside person box
|
|
in_person = False
|
|
person_bbox = None
|
|
for person in person_boxes:
|
|
px1, py1, px2, py2 = person['bbox']
|
|
# Check if segment center is in person box
|
|
center_x = x + w/2
|
|
center_y = y + h/2
|
|
if px1 <= center_x <= px2 and py1 <= center_y <= py2:
|
|
in_person = True
|
|
person_bbox = person['bbox']
|
|
break
|
|
|
|
if in_person and person_bbox:
|
|
px1, py1, px2, py2 = person_bbox
|
|
person_height = py2 - py1
|
|
|
|
# Relative position in person
|
|
rel_y = (y + h/2 - py1) / person_height
|
|
aspect_ratio = w / h if h > 0 else 1
|
|
area_ratio = area / ((px2-px1) * person_height)
|
|
|
|
# Analyze color for hair detection
|
|
segment_pixels = image[mask]
|
|
if len(segment_pixels) > 0:
|
|
mean_color = np.mean(segment_pixels, axis=0)
|
|
# Check for hair-like colors (browns, blacks, blondes)
|
|
is_hair_color = self._is_hair_color(mean_color)
|
|
else:
|
|
is_hair_color = False
|
|
|
|
# Classification logic
|
|
if rel_y < 0.2: # Top 20% of person
|
|
if is_hair_color or y < py1 + person_height * 0.15:
|
|
return 'hair'
|
|
elif aspect_ratio > 0.6 and aspect_ratio < 1.6:
|
|
return 'head_face'
|
|
else:
|
|
return 'hair' # Default top to hair
|
|
|
|
elif rel_y < 0.3: # Upper head/neck region (expanded)
|
|
if area_ratio > 0.03: # Lower threshold
|
|
return 'head_face'
|
|
else:
|
|
return 'accessory' # Could be glasses, earrings
|
|
|
|
elif rel_y < 0.55: # Upper body
|
|
if area_ratio > 0.15:
|
|
return 'shirt' # or jacket
|
|
elif area_ratio > 0.08:
|
|
return 'jacket'
|
|
else:
|
|
return 'accessory'
|
|
|
|
elif rel_y < 0.85: # Lower body
|
|
if area_ratio > 0.15:
|
|
return 'pants' # or skirt
|
|
else:
|
|
return 'accessory'
|
|
|
|
else: # Feet region
|
|
return 'shoes'
|
|
|
|
else:
|
|
# Not in person - it's a scene object
|
|
return 'object'
|
|
|
|
def _is_hair_color(self, color: np.ndarray) -> bool:
|
|
"""
|
|
Check if color is likely hair color
|
|
"""
|
|
r, g, b = color
|
|
|
|
# Brown hair (most common)
|
|
if 20 < r < 150 and 10 < g < 100 and 5 < b < 80:
|
|
return True
|
|
|
|
# Black hair
|
|
if r < 60 and g < 60 and b < 60:
|
|
return True
|
|
|
|
# Blonde hair
|
|
if 150 < r < 250 and 120 < g < 220 and 60 < b < 180:
|
|
return True
|
|
|
|
# Red/auburn hair
|
|
if 100 < r < 200 and 30 < g < 120 and 10 < b < 80 and r > g * 1.3:
|
|
return True
|
|
|
|
return False
|
|
|
|
def _merge_nearby_segments(self, mask: np.ndarray, all_masks: list,
|
|
segment_type: str, bbox: tuple) -> np.ndarray:
|
|
"""
|
|
Merge nearby segments of similar type
|
|
"""
|
|
x, y, w, h = bbox
|
|
merged_mask = mask.copy()
|
|
|
|
# Look for nearby segments to merge
|
|
for other_mask_data in all_masks:
|
|
other_mask = other_mask_data['segmentation']
|
|
other_bbox = other_mask_data['bbox']
|
|
ox, oy, ow, oh = other_bbox
|
|
|
|
# Check if nearby (within 20 pixels)
|
|
x_dist = min(abs(x - (ox + ow)), abs((x + w) - ox))
|
|
y_dist = min(abs(y - (oy + oh)), abs((y + h) - oy))
|
|
|
|
if x_dist < 20 and y_dist < 20:
|
|
# Check if similar type (would be classified the same)
|
|
if segment_type in ['hair', 'head_face']:
|
|
# Merge if touching or very close
|
|
dilated_mask = ndimage.binary_dilation(merged_mask, iterations=10)
|
|
if np.any(dilated_mask & other_mask):
|
|
merged_mask |= other_mask
|
|
|
|
return merged_mask
|
|
|
|
def save_practical_segments(self, results: Dict, output_dir: str = 'practical_output'):
|
|
"""
|
|
Save practical segments organized by category
|
|
"""
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(exist_ok=True)
|
|
|
|
base_name = Path(results['image_path']).stem
|
|
|
|
# Create directories
|
|
dirs = {}
|
|
for category in results['segments'].keys():
|
|
if results['segments'][category]: # Only create if has items
|
|
dirs[category] = output_path / category
|
|
dirs[category].mkdir(exist_ok=True)
|
|
|
|
# Also create combined directory
|
|
all_dir = output_path / 'all_segments'
|
|
all_dir.mkdir(exist_ok=True)
|
|
|
|
print(f"\nSaving practical segments...")
|
|
|
|
# Save each category
|
|
segment_id = 0
|
|
for category, segments in results['segments'].items():
|
|
for i, segment in enumerate(segments):
|
|
if segment['isolated_cropped'].size == 0:
|
|
continue
|
|
|
|
# Create transparent PNG
|
|
cropped_rgba = np.zeros((*segment['isolated_cropped'].shape[:2], 4), dtype=np.uint8)
|
|
cropped_rgba[:, :, :3] = segment['isolated_cropped']
|
|
cropped_rgba[:, :, 3] = segment['mask_cropped'] * 255
|
|
|
|
# Determine filename
|
|
if category == 'clothing' and 'clothing_type' in segment:
|
|
item_name = segment['clothing_type']
|
|
elif category == 'scene_objects' and 'type' in segment:
|
|
item_name = segment['type']
|
|
else:
|
|
item_name = category.rstrip('s') # Remove plural
|
|
|
|
# Save to category folder
|
|
if category in dirs:
|
|
cat_path = dirs[category] / f"{base_name}_{item_name}_{i:02d}.png"
|
|
Image.fromarray(cropped_rgba).save(str(cat_path))
|
|
|
|
# Save to all folder
|
|
all_path = all_dir / f"{base_name}_{segment_id:03d}_{item_name}.png"
|
|
Image.fromarray(cropped_rgba).save(str(all_path))
|
|
|
|
segment_id += 1
|
|
|
|
# Create visualization
|
|
self._create_visualization(results, output_path / f"{base_name}_practical.jpg")
|
|
|
|
# Save summary
|
|
summary_path = output_path / f"{base_name}_summary.txt"
|
|
with open(summary_path, 'w') as f:
|
|
f.write(f"Practical Segmentation Results\n")
|
|
f.write(f"{'='*40}\n")
|
|
f.write(f"Image: {results['image_path']}\n")
|
|
f.write(f"Size: {results['image_size'][0]}x{results['image_size'][1]}\n\n")
|
|
|
|
total = sum(len(segs) for segs in results['segments'].values())
|
|
f.write(f"Total segments: {total}\n\n")
|
|
|
|
f.write("Segments by category:\n")
|
|
for category, segments in results['segments'].items():
|
|
if segments:
|
|
f.write(f"\n{category.upper()}:\n")
|
|
for i, seg in enumerate(segments):
|
|
if 'clothing_type' in seg:
|
|
name = seg['clothing_type']
|
|
elif 'type' in seg:
|
|
name = seg['type']
|
|
else:
|
|
name = category.rstrip('s')
|
|
f.write(f" [{i}] {name} - {seg['area']} pixels\n")
|
|
|
|
print(f"Results saved to: {output_path}")
|
|
|
|
# Print summary
|
|
print("\nSummary:")
|
|
for category, segments in results['segments'].items():
|
|
if segments:
|
|
print(f" {category}: {len(segments)} items")
|
|
|
|
def _create_visualization(self, results: Dict, output_path: Path):
|
|
"""
|
|
Create clean visualization
|
|
"""
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
|
|
fig, axes = plt.subplots(2, 2, figsize=(12, 12))
|
|
|
|
# Original
|
|
axes[0, 0].imshow(results['original_image'])
|
|
axes[0, 0].set_title('Original Image')
|
|
axes[0, 0].axis('off')
|
|
|
|
# All segments with labels
|
|
axes[0, 1].imshow(results['original_image'])
|
|
axes[0, 1].set_title('Practical Segments')
|
|
|
|
# Define colors for each category
|
|
category_colors = {
|
|
'hair': 'purple',
|
|
'head_face': 'yellow',
|
|
'body': 'cyan',
|
|
'clothing': 'green',
|
|
'accessories': 'orange',
|
|
'scene_objects': 'red'
|
|
}
|
|
|
|
for category, segments in results['segments'].items():
|
|
color = category_colors.get(category, 'blue')
|
|
for seg in segments:
|
|
x1, y1, x2, y2 = seg['bbox']
|
|
rect = patches.Rectangle((x1, y1), x2-x1, y2-y1,
|
|
linewidth=2, edgecolor=color,
|
|
facecolor='none')
|
|
axes[0, 1].add_patch(rect)
|
|
|
|
# Add label
|
|
label = seg.get('clothing_type', seg.get('type', category.rstrip('s')))
|
|
axes[0, 1].text(x1, y1-2, label, fontsize=8, color=color,
|
|
bbox=dict(boxstyle='round,pad=0.2',
|
|
facecolor='white', alpha=0.8))
|
|
axes[0, 1].axis('off')
|
|
|
|
# Mask overlay
|
|
mask_overlay = np.zeros((*results['original_image'].shape[:2], 3))
|
|
|
|
segment_idx = 0
|
|
for category, segments in results['segments'].items():
|
|
for seg in segments:
|
|
color = plt.cm.rainbow(segment_idx / 20)[:3]
|
|
mask_overlay[seg['mask'] == 1] = color
|
|
segment_idx += 1
|
|
|
|
axes[1, 0].imshow(mask_overlay)
|
|
axes[1, 0].set_title('Segment Masks')
|
|
axes[1, 0].axis('off')
|
|
|
|
# Overlay
|
|
axes[1, 1].imshow(results['original_image'])
|
|
axes[1, 1].imshow(mask_overlay, alpha=0.4)
|
|
axes[1, 1].set_title('Segments Overlay')
|
|
axes[1, 1].axis('off')
|
|
|
|
plt.tight_layout()
|
|
plt.savefig(output_path, dpi=150, bbox_inches='tight')
|
|
plt.close()
|
|
|
|
def main():
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description='Practical segmentation - no tiny body parts!')
|
|
parser.add_argument('images', nargs='+', help='Image file(s) to process')
|
|
parser.add_argument('-o', '--output', default='practical_output', help='Output directory')
|
|
parser.add_argument('--no-merge', action='store_true', help="Don't merge body parts")
|
|
parser.add_argument('-d', '--device', default='auto', help='Device (cpu, cuda, auto)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Initialize
|
|
segmenter = PracticalSegmenter(device=args.device)
|
|
|
|
# Process images
|
|
for image_path in args.images:
|
|
if Path(image_path).is_file():
|
|
try:
|
|
results = segmenter.segment_practical(
|
|
image_path,
|
|
merge_body_parts=not args.no_merge
|
|
)
|
|
segmenter.save_practical_segments(results, args.output)
|
|
print()
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
print("Done!")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|