File size: 4,675 Bytes
58343f2 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
# sample_inference.py
import argparse
import torch
from ultralytics import YOLO
import cv2
import numpy as np
import json
from PIL import Image
def main():
parser = argparse.ArgumentParser(description='Run person segmentation with YOLO12l-seg model')
parser.add_argument('--model', type=str, default='yolo12l-person-seg.pt', help='Model path')
parser.add_argument('--image', type=str, required=True, help='Image path for inference')
parser.add_argument('--output', type=str, default='output.jpg', help='Output visualization image path')
parser.add_argument('--json', type=str, default='detections.json', help='JSON output file for detection data')
parser.add_argument('--conf', type=float, default=0.5, help='Confidence threshold')
args = parser.parse_args()
# Load the model
model = YOLO(args.model)
# Move to appropriate device if available
if torch.cuda.is_available():
print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
model.to('cuda')
device = 'cuda'
use_half = True
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
print("Using Apple Silicon MPS")
model.to('mps')
device = 'mps'
use_half = False
else:
print("Using CPU")
device = None
use_half = False
# Load and check input image
try:
img = Image.open(args.image)
img_width, img_height = img.size
print(f"Image dimensions: {img_width}x{img_height}")
except Exception as e:
print(f"Error opening image: {e}")
return
# Run inference
if device == 'cuda':
results = model(args.image, classes=0, conf=args.conf, device=device, half=use_half)
elif device == 'mps':
results = model(args.image, classes=0, conf=args.conf, device=device)
else:
results = model(args.image, classes=0, conf=args.conf)
# Process results
detections = []
visualization_img = cv2.imread(args.image)
for result in results:
masks = result.masks
boxes = result.boxes
if boxes is None or len(boxes) == 0:
print("No people detected in the image")
return
person_count = len(boxes)
print(f"Detected {person_count} people")
# Visualize and extract data
if masks is not None:
for i, (mask, box) in enumerate(zip(masks.xy, boxes)):
confidence = float(box.conf[0])
x1, y1, x2, y2 = map(int, box.xyxy[0])
# Extract mask points
polygon_points = mask.tolist()
# Calculate percentages of image dimensions
x_coords = [point[0] for point in polygon_points]
y_coords = [point[1] for point in polygon_points]
min_x, max_x = min(x_coords), max(x_coords)
min_y, max_y = min(y_coords), max(y_coords)
width_pct = (max_x - min_x) / img_width
height_pct = (max_y - min_y) / img_height
# Create detection record
detection = {
"id": i,
"confidence": confidence,
"box": [x1, y1, x2, y2],
"points": polygon_points,
"width_pct": width_pct,
"height_pct": height_pct,
}
detections.append(detection)
# Draw bounding box
cv2.rectangle(visualization_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(visualization_img, f'Person: {confidence:.2f}', (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# Draw segmentation mask
color_mask = np.zeros_like(visualization_img, dtype=np.uint8)
mask_points = np.array(polygon_points, dtype=np.int32)
cv2.fillPoly(color_mask, [mask_points], (0, 0, 255))
# Blend the mask with the original image
visualization_img = cv2.addWeighted(visualization_img, 1.0, color_mask, 0.5, 0)
# Save visualization
cv2.imwrite(args.output, visualization_img)
print(f"Visualization saved to {args.output}")
# Save detection data to JSON
with open(args.json, 'w') as f:
json.dump({
"person_count": person_count,
"detections": detections
}, f, indent=4)
print(f"Detection data saved to {args.json}")
if __name__ == "__main__":
main() |