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
376 lines
14 KiB
Python
376 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Segment EVERYTHING in an image - every object, clothing item, and element
|
|
Uses SAM (Segment Anything Model) for comprehensive segmentation
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
import supervision as sv
|
|
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
|
|
from typing import List, Dict, Any
|
|
import requests
|
|
from tqdm import tqdm
|
|
|
|
class EverythingSegmenter:
|
|
def __init__(self, model_type='vit_b', device='auto'):
|
|
"""
|
|
Initialize SAM for segmenting everything
|
|
|
|
Args:
|
|
model_type: SAM model type ('vit_b', 'vit_l', 'vit_h')
|
|
device: Device to use
|
|
"""
|
|
# Auto-detect device
|
|
if device == 'auto':
|
|
if torch.cuda.is_available():
|
|
self.device = 'cuda'
|
|
print("Using CUDA (GPU)")
|
|
elif torch.backends.mps.is_available():
|
|
self.device = 'cpu' # SAM doesn't work well with MPS yet
|
|
print("Using CPU (SAM doesn't support MPS yet)")
|
|
else:
|
|
self.device = 'cpu'
|
|
print("Using CPU")
|
|
else:
|
|
self.device = device
|
|
|
|
# Model URLs
|
|
model_urls = {
|
|
'vit_b': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
|
|
'vit_l': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
|
|
'vit_h': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth'
|
|
}
|
|
|
|
# Download model if needed
|
|
model_path = Path(f'sam_{model_type}.pth')
|
|
if not model_path.exists():
|
|
print(f"Downloading SAM {model_type} model...")
|
|
self._download_model(model_urls[model_type], model_path)
|
|
|
|
# Load SAM model
|
|
print(f"Loading SAM {model_type} model...")
|
|
self.sam = sam_model_registry[model_type](checkpoint=str(model_path))
|
|
self.sam.to(device=self.device)
|
|
|
|
# Create mask generator with optimized parameters
|
|
self.mask_generator = SamAutomaticMaskGenerator(
|
|
self.sam,
|
|
points_per_side=32, # More points = more masks
|
|
pred_iou_thresh=0.86, # Quality threshold
|
|
stability_score_thresh=0.92, # Stability threshold
|
|
crop_n_layers=1, # Use crops for better small object detection
|
|
crop_n_points_downscale_factor=2,
|
|
min_mask_region_area=100, # Minimum mask size in pixels
|
|
)
|
|
|
|
print("Model loaded successfully!")
|
|
|
|
def _download_model(self, url: str, path: Path):
|
|
"""Download SAM model weights"""
|
|
response = requests.get(url, stream=True)
|
|
total_size = int(response.headers.get('content-length', 0))
|
|
|
|
with open(path, 'wb') as f:
|
|
with tqdm(total=total_size, unit='B', unit_scale=True, desc='Downloading') as pbar:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
pbar.update(len(chunk))
|
|
|
|
def segment_everything(self, image_path: str, min_area: int = 100) -> Dict[str, Any]:
|
|
"""
|
|
Segment everything in an image
|
|
|
|
Args:
|
|
image_path: Path to input image
|
|
min_area: Minimum area for a mask to be kept
|
|
|
|
Returns:
|
|
Dictionary containing all masks and metadata
|
|
"""
|
|
# 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"Processing {Path(image_path).name} ({width}x{height})...")
|
|
print("Generating masks for everything in the image...")
|
|
|
|
# Generate masks
|
|
masks = self.mask_generator.generate(image_rgb)
|
|
|
|
# Sort masks by area (largest first)
|
|
masks = sorted(masks, key=lambda x: x['area'], reverse=True)
|
|
|
|
# Filter small masks
|
|
masks = [m for m in masks if m['area'] >= min_area]
|
|
|
|
print(f"Found {len(masks)} distinct segments")
|
|
|
|
# Process masks
|
|
results = {
|
|
'original_image': image_rgb,
|
|
'image_path': image_path,
|
|
'image_size': (width, height),
|
|
'segments': []
|
|
}
|
|
|
|
for i, mask_data in enumerate(masks):
|
|
# Get mask
|
|
mask = mask_data['segmentation'].astype(np.uint8)
|
|
|
|
# Get bounding box
|
|
bbox = mask_data['bbox'] # x, y, w, h format
|
|
x, y, w, h = [int(v) for v in bbox]
|
|
|
|
# Extract isolated object
|
|
isolated_full = image_rgb.copy()
|
|
isolated_full[mask == 0] = 0
|
|
|
|
# Crop to bounding box
|
|
isolated_cropped = isolated_full[y:y+h, x:x+w]
|
|
mask_cropped = mask[y:y+h, x:x+w]
|
|
|
|
# Estimate what this might be based on position and size
|
|
segment_type = self._estimate_segment_type(mask_data, image.shape, i)
|
|
|
|
segment_info = {
|
|
'id': i,
|
|
'type': segment_type,
|
|
'area': mask_data['area'],
|
|
'bbox': (x, y, x+w, y+h),
|
|
'mask': mask,
|
|
'isolated_full': isolated_full,
|
|
'isolated_cropped': isolated_cropped,
|
|
'mask_cropped': mask_cropped,
|
|
'stability_score': mask_data['stability_score'],
|
|
'predicted_iou': mask_data['predicted_iou']
|
|
}
|
|
results['segments'].append(segment_info)
|
|
|
|
return results
|
|
|
|
def _estimate_segment_type(self, mask_data: Dict, image_shape: tuple, index: int) -> str:
|
|
"""Estimate what type of segment this might be based on position and size"""
|
|
bbox = mask_data['bbox']
|
|
x, y, w, h = bbox
|
|
area_ratio = mask_data['area'] / (image_shape[0] * image_shape[1])
|
|
aspect_ratio = w / h if h > 0 else 1
|
|
|
|
# Position in image (normalized)
|
|
center_x = (x + w/2) / image_shape[1]
|
|
center_y = (y + h/2) / image_shape[0]
|
|
|
|
# Simple heuristics to guess segment type
|
|
if area_ratio > 0.4:
|
|
return "background"
|
|
elif area_ratio > 0.15:
|
|
if center_y < 0.5 and aspect_ratio > 0.5 and aspect_ratio < 2:
|
|
return "person_upper"
|
|
elif center_y > 0.5:
|
|
return "person_lower"
|
|
else:
|
|
return "large_object"
|
|
elif y < image_shape[0] * 0.3: # Top third
|
|
if aspect_ratio > 0.8 and aspect_ratio < 1.5:
|
|
return "head_item" # Could be head, hat, hair
|
|
else:
|
|
return "upper_item"
|
|
elif y > image_shape[0] * 0.7: # Bottom third
|
|
return "lower_item" # Could be shoes, pants bottom
|
|
else: # Middle
|
|
if w > h:
|
|
return "horizontal_item"
|
|
else:
|
|
return "vertical_item"
|
|
|
|
def save_everything(self, results: Dict, output_dir: str = 'everything_output'):
|
|
"""
|
|
Save all segmented items
|
|
|
|
Args:
|
|
results: Results from segment_everything
|
|
output_dir: Output directory
|
|
"""
|
|
# Create output directories
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(exist_ok=True)
|
|
|
|
base_name = Path(results['image_path']).stem
|
|
|
|
# Create subdirectories by type
|
|
dirs = {
|
|
'all_masks': output_path / 'all_masks',
|
|
'isolated': output_path / 'isolated',
|
|
'by_type': output_path / 'by_type',
|
|
'transparent': output_path / 'transparent',
|
|
'visualization': output_path / 'visualization'
|
|
}
|
|
|
|
for dir_path in dirs.values():
|
|
dir_path.mkdir(exist_ok=True)
|
|
|
|
print(f"\nSaving {len(results['segments'])} segments...")
|
|
|
|
# Group segments by type
|
|
segments_by_type = {}
|
|
for segment in results['segments']:
|
|
seg_type = segment['type']
|
|
if seg_type not in segments_by_type:
|
|
segments_by_type[seg_type] = []
|
|
segments_by_type[seg_type].append(segment)
|
|
|
|
# Save individual segments
|
|
for segment in results['segments']:
|
|
seg_id = segment['id']
|
|
seg_type = segment['type']
|
|
|
|
# Create type directory
|
|
type_dir = dirs['by_type'] / seg_type
|
|
type_dir.mkdir(exist_ok=True)
|
|
|
|
# Save mask
|
|
mask_path = dirs['all_masks'] / f"{base_name}_{seg_id:03d}_{seg_type}_mask.png"
|
|
cv2.imwrite(str(mask_path), segment['mask'] * 255)
|
|
|
|
# Save isolated (full size)
|
|
isolated_path = dirs['isolated'] / f"{base_name}_{seg_id:03d}_{seg_type}.jpg"
|
|
cv2.imwrite(str(isolated_path), cv2.cvtColor(segment['isolated_full'], cv2.COLOR_RGB2BGR))
|
|
|
|
# Save transparent cropped
|
|
if segment['isolated_cropped'].size > 0:
|
|
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
|
|
|
|
trans_path = dirs['transparent'] / f"{base_name}_{seg_id:03d}_{seg_type}.png"
|
|
Image.fromarray(cropped_rgba).save(str(trans_path))
|
|
|
|
# Also save in type directory
|
|
type_trans_path = type_dir / f"{base_name}_{seg_id:03d}.png"
|
|
Image.fromarray(cropped_rgba).save(str(type_trans_path))
|
|
|
|
# Create visualization
|
|
self._create_visualization(results, dirs['visualization'] / f"{base_name}_overview.jpg")
|
|
|
|
# Save summary
|
|
summary_path = output_path / f"{base_name}_summary.txt"
|
|
with open(summary_path, 'w') as f:
|
|
f.write(f"Image: {results['image_path']}\n")
|
|
f.write(f"Size: {results['image_size'][0]}x{results['image_size'][1]}\n")
|
|
f.write(f"Total segments: {len(results['segments'])}\n\n")
|
|
|
|
f.write("Segments by type:\n")
|
|
for seg_type, segments in segments_by_type.items():
|
|
f.write(f" {seg_type}: {len(segments)} segments\n")
|
|
|
|
f.write("\nAll segments:\n")
|
|
for segment in results['segments']:
|
|
f.write(f" [{segment['id']:3d}] {segment['type']:15s} - Area: {segment['area']:6d} pixels\n")
|
|
|
|
print(f"Results saved to: {output_path}")
|
|
print(f"\nSegments by type:")
|
|
for seg_type, segments in segments_by_type.items():
|
|
print(f" {seg_type}: {len(segments)} segments")
|
|
|
|
def _create_visualization(self, results: Dict, output_path: Path):
|
|
"""Create visualization showing all segments"""
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
|
|
fig, axes = plt.subplots(2, 2, figsize=(12, 12))
|
|
|
|
# Original image
|
|
axes[0, 0].imshow(results['original_image'])
|
|
axes[0, 0].set_title('Original Image')
|
|
axes[0, 0].axis('off')
|
|
|
|
# All masks with different colors
|
|
mask_overlay = np.zeros((*results['original_image'].shape[:2], 3))
|
|
colors = plt.cm.rainbow(np.linspace(0, 1, len(results['segments'])))
|
|
|
|
for i, segment in enumerate(results['segments']):
|
|
mask = segment['mask']
|
|
color = colors[i][:3]
|
|
mask_overlay[mask == 1] = color
|
|
|
|
axes[0, 1].imshow(mask_overlay)
|
|
axes[0, 1].set_title(f'All {len(results["segments"])} Segments')
|
|
axes[0, 1].axis('off')
|
|
|
|
# Overlay on original
|
|
axes[1, 0].imshow(results['original_image'])
|
|
axes[1, 0].imshow(mask_overlay, alpha=0.5)
|
|
axes[1, 0].set_title('Segments Overlay')
|
|
axes[1, 0].axis('off')
|
|
|
|
# Bounding boxes with labels
|
|
axes[1, 1].imshow(results['original_image'])
|
|
|
|
# Group by type for coloring
|
|
type_colors = {}
|
|
unique_types = list(set(s['type'] for s in results['segments']))
|
|
type_color_map = plt.cm.tab10(np.linspace(0, 1, len(unique_types)))
|
|
for i, t in enumerate(unique_types):
|
|
type_colors[t] = type_color_map[i]
|
|
|
|
for segment in results['segments'][:20]: # Show first 20 to avoid clutter
|
|
x1, y1, x2, y2 = segment['bbox']
|
|
rect = patches.Rectangle((x1, y1), x2-x1, y2-y1,
|
|
linewidth=1, edgecolor=type_colors[segment['type']],
|
|
facecolor='none')
|
|
axes[1, 1].add_patch(rect)
|
|
|
|
# Add label
|
|
axes[1, 1].text(x1, y1-2, f"{segment['id']}:{segment['type'][:8]}",
|
|
fontsize=6, color=type_colors[segment['type']],
|
|
bbox=dict(boxstyle='round,pad=0.1', facecolor='white', alpha=0.7))
|
|
|
|
axes[1, 1].set_title('Segment Types (first 20)')
|
|
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='Segment EVERYTHING in an image')
|
|
parser.add_argument('images', nargs='+', help='Image file(s) to process')
|
|
parser.add_argument('-o', '--output', default='everything_output', help='Output directory')
|
|
parser.add_argument('-m', '--model', default='vit_b', choices=['vit_b', 'vit_l', 'vit_h'],
|
|
help='SAM model size (vit_b=base, vit_l=large, vit_h=huge)')
|
|
parser.add_argument('--min-area', type=int, default=100,
|
|
help='Minimum segment area in pixels')
|
|
parser.add_argument('-d', '--device', default='auto', help='Device (cpu, cuda, auto)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Initialize segmenter
|
|
segmenter = EverythingSegmenter(model_type=args.model, device=args.device)
|
|
|
|
# Process images
|
|
for image_path in args.images:
|
|
if Path(image_path).is_file():
|
|
try:
|
|
results = segmenter.segment_everything(image_path, min_area=args.min_area)
|
|
segmenter.save_everything(results, args.output)
|
|
print()
|
|
except Exception as e:
|
|
print(f"Error processing {image_path}: {e}")
|
|
continue
|
|
|
|
print("Done!")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|