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
327 lines
12 KiB
Python
327 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Object Segmentation Script using YOLOv8
|
|
Segments all objects in an image and outputs masks or isolated objects
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
from pathlib import Path
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
from ultralytics import YOLO
|
|
import matplotlib.pyplot as plt
|
|
from typing import List, Tuple, Optional
|
|
|
|
class ObjectSegmenter:
|
|
def __init__(self, model_name='yolov8m-seg.pt', device='auto'):
|
|
"""
|
|
Initialize the segmentation model
|
|
|
|
Args:
|
|
model_name: YOLOv8 segmentation model to use
|
|
Options: yolov8n-seg, yolov8s-seg, yolov8m-seg, yolov8l-seg, yolov8x-seg
|
|
device: Device to use ('cpu', 'mps' for Mac, 'cuda' for GPU, or 'auto')
|
|
"""
|
|
# 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 = 'mps'
|
|
print("Using MPS (Apple Silicon)")
|
|
else:
|
|
self.device = 'cpu'
|
|
print("Using CPU")
|
|
else:
|
|
self.device = device
|
|
|
|
# Load YOLO model
|
|
print(f"Loading {model_name} model...")
|
|
self.model = YOLO(model_name)
|
|
self.model.to(self.device)
|
|
print("Model loaded successfully!")
|
|
|
|
# Get class names
|
|
self.class_names = self.model.names
|
|
|
|
def segment_image(self, image_path: str, conf_threshold: float = 0.25) -> dict:
|
|
"""
|
|
Segment all objects in an image
|
|
|
|
Args:
|
|
image_path: Path to input image
|
|
conf_threshold: Confidence threshold for detections
|
|
|
|
Returns:
|
|
Dictionary containing results
|
|
"""
|
|
# Read image
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
raise ValueError(f"Could not read image: {image_path}")
|
|
|
|
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
height, width = img.shape[:2]
|
|
|
|
# Run inference
|
|
results = self.model(img_rgb, conf=conf_threshold)
|
|
|
|
# Process results
|
|
output = {
|
|
'original_image': img_rgb,
|
|
'image_path': image_path,
|
|
'image_size': (width, height),
|
|
'objects': []
|
|
}
|
|
|
|
if len(results) > 0 and results[0].masks is not None:
|
|
result = results[0]
|
|
masks = result.masks.data.cpu().numpy()
|
|
boxes = result.boxes.xyxy.cpu().numpy()
|
|
classes = result.boxes.cls.cpu().numpy().astype(int)
|
|
confidences = result.boxes.conf.cpu().numpy()
|
|
|
|
for i in range(len(masks)):
|
|
# Get mask
|
|
mask = masks[i]
|
|
|
|
# Resize mask to original image size
|
|
mask_resized = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST)
|
|
mask_binary = (mask_resized > 0.5).astype(np.uint8)
|
|
|
|
# Extract object using mask
|
|
isolated_object = img_rgb.copy()
|
|
isolated_object[mask_binary == 0] = 0 # Black background
|
|
|
|
# Get bounding box
|
|
x1, y1, x2, y2 = boxes[i].astype(int)
|
|
cropped_object = isolated_object[y1:y2, x1:x2]
|
|
cropped_mask = mask_binary[y1:y2, x1:x2]
|
|
|
|
# Store object info
|
|
obj_info = {
|
|
'id': i,
|
|
'class_name': self.class_names[classes[i]],
|
|
'class_id': classes[i],
|
|
'confidence': confidences[i],
|
|
'bbox': (x1, y1, x2, y2),
|
|
'mask': mask_binary,
|
|
'mask_resized': mask_resized,
|
|
'isolated_full': isolated_object,
|
|
'isolated_cropped': cropped_object,
|
|
'cropped_mask': cropped_mask
|
|
}
|
|
output['objects'].append(obj_info)
|
|
|
|
return output
|
|
|
|
def save_results(self, results: dict, output_dir: str = 'output'):
|
|
"""
|
|
Save segmentation results to files
|
|
|
|
Args:
|
|
results: Results dictionary from segment_image
|
|
output_dir: Directory to save outputs
|
|
"""
|
|
# Create output directory
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(exist_ok=True)
|
|
|
|
# Get base name from input image
|
|
base_name = Path(results['image_path']).stem
|
|
|
|
# Create subdirectories
|
|
masks_dir = output_path / 'masks'
|
|
isolated_dir = output_path / 'isolated'
|
|
cropped_dir = output_path / 'cropped'
|
|
composite_dir = output_path / 'composite'
|
|
|
|
for dir_path in [masks_dir, isolated_dir, cropped_dir, composite_dir]:
|
|
dir_path.mkdir(exist_ok=True)
|
|
|
|
# Save individual objects
|
|
print(f"\nFound {len(results['objects'])} objects in {base_name}")
|
|
for obj in results['objects']:
|
|
obj_id = obj['id']
|
|
class_name = obj['class_name'].replace(' ', '_')
|
|
confidence = obj['confidence']
|
|
|
|
# File names
|
|
prefix = f"{base_name}_obj{obj_id:02d}_{class_name}_{confidence:.2f}"
|
|
|
|
# Save mask
|
|
mask_path = masks_dir / f"{prefix}_mask.png"
|
|
cv2.imwrite(str(mask_path), obj['mask'] * 255)
|
|
|
|
# Save isolated object (full image with background removed)
|
|
isolated_path = isolated_dir / f"{prefix}_isolated.png"
|
|
cv2.imwrite(str(isolated_path), cv2.cvtColor(obj['isolated_full'], cv2.COLOR_RGB2BGR))
|
|
|
|
# Save cropped object
|
|
if obj['isolated_cropped'].size > 0:
|
|
cropped_path = cropped_dir / f"{prefix}_cropped.png"
|
|
cv2.imwrite(str(cropped_path), cv2.cvtColor(obj['isolated_cropped'], cv2.COLOR_RGB2BGR))
|
|
|
|
# Save cropped with transparency
|
|
cropped_rgba = np.zeros((*obj['isolated_cropped'].shape[:2], 4), dtype=np.uint8)
|
|
cropped_rgba[:, :, :3] = obj['isolated_cropped']
|
|
cropped_rgba[:, :, 3] = obj['cropped_mask'] * 255
|
|
|
|
cropped_alpha_path = cropped_dir / f"{prefix}_cropped_alpha.png"
|
|
Image.fromarray(cropped_rgba).save(str(cropped_alpha_path))
|
|
|
|
print(f" [{obj_id:2d}] {class_name}: {confidence:.2%} - Saved")
|
|
|
|
# Create and save composite visualization
|
|
self._save_composite(results, composite_dir / f"{base_name}_composite.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"Objects found: {len(results['objects'])}\n\n")
|
|
|
|
for obj in results['objects']:
|
|
f.write(f"Object {obj['id']:2d}: {obj['class_name']}\n")
|
|
f.write(f" Confidence: {obj['confidence']:.2%}\n")
|
|
f.write(f" Bounding box: {obj['bbox']}\n\n")
|
|
|
|
print(f"\nResults saved to: {output_path}")
|
|
|
|
def _save_composite(self, results: dict, output_path: Path):
|
|
"""Create and save a composite visualization"""
|
|
n_objects = len(results['objects'])
|
|
if n_objects == 0:
|
|
print("No objects detected, skipping composite")
|
|
return
|
|
|
|
# Calculate grid size
|
|
cols = min(4, n_objects)
|
|
rows = (n_objects + cols - 1) // cols
|
|
|
|
# Create figure
|
|
fig_height = max(8, rows * 3)
|
|
fig, axes = plt.subplots(rows + 1, cols, figsize=(cols * 3, fig_height))
|
|
|
|
# Ensure axes is always 2D
|
|
if (rows + 1) == 1 and cols == 1:
|
|
axes = np.array([[axes]])
|
|
elif (rows + 1) == 1:
|
|
axes = axes.reshape(1, -1)
|
|
elif cols == 1:
|
|
axes = axes.reshape(-1, 1)
|
|
|
|
# Show original with all masks
|
|
ax = axes[0, 0]
|
|
ax.imshow(results['original_image'])
|
|
ax.set_title('Original + All Masks')
|
|
ax.axis('off')
|
|
|
|
# Overlay all masks with different colors
|
|
colors = plt.cm.rainbow(np.linspace(0, 1, n_objects))
|
|
for i, obj in enumerate(results['objects']):
|
|
mask = obj['mask']
|
|
colored_mask = np.zeros((*mask.shape, 4))
|
|
colored_mask[mask == 1] = colors[i]
|
|
colored_mask[mask == 1, 3] = 0.4 # Alpha
|
|
ax.imshow(colored_mask)
|
|
|
|
# Hide other cells in first row
|
|
for j in range(1, cols):
|
|
axes[0, j].axis('off')
|
|
|
|
# Show individual objects
|
|
for i, obj in enumerate(results['objects']):
|
|
row = (i // cols) + 1
|
|
col = i % cols
|
|
|
|
ax = axes[row, col]
|
|
|
|
# Show cropped object
|
|
if obj['isolated_cropped'].size > 0:
|
|
ax.imshow(obj['isolated_cropped'])
|
|
title = f"{obj['class_name']}\n{obj['confidence']:.1%}"
|
|
ax.set_title(title, fontsize=9)
|
|
ax.axis('off')
|
|
|
|
# Hide empty cells
|
|
for i in range(n_objects, rows * cols):
|
|
row = (i // cols) + 1
|
|
col = i % cols
|
|
ax = axes[row, col]
|
|
ax.axis('off')
|
|
|
|
plt.tight_layout()
|
|
plt.savefig(output_path, dpi=150, bbox_inches='tight')
|
|
plt.close()
|
|
|
|
def process_batch(self, image_paths: List[str], output_dir: str = 'output',
|
|
conf_threshold: float = 0.25):
|
|
"""
|
|
Process multiple images
|
|
|
|
Args:
|
|
image_paths: List of image paths
|
|
output_dir: Output directory
|
|
conf_threshold: Confidence threshold
|
|
"""
|
|
print(f"\nProcessing {len(image_paths)} images...")
|
|
|
|
for i, image_path in enumerate(image_paths, 1):
|
|
print(f"\n[{i}/{len(image_paths)}] Processing: {image_path}")
|
|
try:
|
|
results = self.segment_image(image_path, conf_threshold)
|
|
self.save_results(results, output_dir)
|
|
except Exception as e:
|
|
print(f" Error processing {image_path}: {e}")
|
|
continue
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Segment objects in images using YOLOv8')
|
|
parser.add_argument('input', nargs='+', help='Input image(s) or directory')
|
|
parser.add_argument('--output', '-o', default='output', help='Output directory')
|
|
parser.add_argument('--model', '-m', default='yolov8m-seg.pt',
|
|
help='Model to use (yolov8n-seg, yolov8s-seg, yolov8m-seg, yolov8l-seg, yolov8x-seg)')
|
|
parser.add_argument('--conf', '-c', type=float, default=0.25,
|
|
help='Confidence threshold (0-1)')
|
|
parser.add_argument('--device', '-d', default='auto',
|
|
help='Device to use (cpu, mps, cuda, or auto)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Collect image paths
|
|
image_paths = []
|
|
for input_path in args.input:
|
|
path = Path(input_path)
|
|
if path.is_file():
|
|
if path.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
|
|
image_paths.append(str(path))
|
|
elif path.is_dir():
|
|
for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff']:
|
|
image_paths.extend([str(p) for p in path.glob(ext)])
|
|
|
|
if not image_paths:
|
|
print("No valid images found!")
|
|
return
|
|
|
|
# Initialize segmenter
|
|
segmenter = ObjectSegmenter(model_name=args.model, device=args.device)
|
|
|
|
# Process images
|
|
if len(image_paths) == 1:
|
|
results = segmenter.segment_image(image_paths[0], args.conf)
|
|
segmenter.save_results(results, args.output)
|
|
else:
|
|
segmenter.process_batch(image_paths, args.output, args.conf)
|
|
|
|
print("\nDone!")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|