File size: 6,561 Bytes
6a5bb7e | 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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | """Snake Active Contour Bbox Refinement
========================================
Classic CV approach: run active contour on image gradient inside each YOLO box,
fit the snake to the goat's actual boundary, then fit a tighter box.
Zero training cost. CPU-based.
"""
import sys, os
import numpy as np
from PIL import Image
from tqdm import tqdm
from scipy import ndimage
PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)
from ultralytics import YOLO
def snake_refine(image, box, n_iter=30, alpha=0.01, beta=0.1, gamma=0.001):
"""Refine a single bounding box using active contour.
Args:
image: PIL Image (full image)
box: [x1, y1, x2, y2] in pixel coords
n_iter: snake iterations
alpha: continuity (smoothness along contour)
beta: curvature (bending penalty)
gamma: step size
Returns:
refined_box: [x1, y1, x2, y2] or None if failed
"""
x1, y1, x2, y2 = box
w, h = x2 - x1, y2 - y1
if w < 8 or h < 8:
return None
# Expand ROI for context
pad = 0.3
rx1 = max(0, int(x1 - w * pad))
ry1 = max(0, int(y1 - h * pad))
rx2 = min(image.size[0], int(x2 + w * pad))
ry2 = min(image.size[1], int(y2 + h * pad))
roi = image.crop((rx1, ry1, rx2, ry2))
roi_gray = np.array(roi.convert('L'), dtype=float)
# Edge map (negative gradient magnitude attracts snake)
gy, gx = np.gradient(roi_gray)
edge = np.sqrt(gx**2 + gy**2)
edge = edge / (edge.max() + 1e-8)
# Initialize snake as ellipse inside the original box (in ROI coords)
roi_h, roi_w = roi_gray.shape
cx_roi = (x1 + x2) / 2 - rx1
cy_roi = (y1 + y2) / 2 - ry1
rw = w * 0.9 # Slightly smaller than YOLO box
rh = h * 0.9
# Create initial contour (ellipse with 40 points)
n_pts = 40
theta = np.linspace(0, 2*np.pi, n_pts)
snake_x = cx_roi + rw/2 * np.cos(theta)
snake_y = cy_roi + rh/2 * np.sin(theta)
# Clip to ROI
snake_x = np.clip(snake_x, 1, roi_w - 2)
snake_y = np.clip(snake_y, 1, roi_h - 2)
# Precompute external energy: negative edge + distance to edge
external = -edge
# Add small inward balloon force
balloon = -0.05
# Snake evolution
for _ in range(n_iter):
# Get snake as integer coords for interpolation
sx = np.clip(snake_x.astype(int), 0, roi_w - 1)
sy = np.clip(snake_y.astype(int), 0, roi_h - 1)
# External force: gradient of edge map at snake points
# Use simple gradient descent on external energy
for i in range(n_pts):
x, y = int(snake_x[i]), int(snake_y[i])
x = max(1, min(roi_w - 2, x))
y = max(1, min(roi_h - 2, y))
# Local gradient of external energy
dx = external[y, min(x+1, roi_w-1)] - external[y, max(x-1, 0)]
dy = external[min(y+1, roi_h-1), x] - external[max(y-1, 0), x]
# Internal force: tension + rigidity
prev_i = (i - 1) % n_pts
next_i = (i + 1) % n_pts
# Tension: pull towards neighbors (alpha)
tension = alpha * (snake_x[prev_i] + snake_x[next_i] - 2*snake_x[i])
# Update
snake_x[i] += gamma * (-dx * 50 + tension) + balloon * (snake_x[i] - cx_roi) * 0.01
snake_y[i] += gamma * (-dy * 50 + tension) + balloon * (snake_y[i] - cy_roi) * 0.01
# Re-parameterize: resample evenly
# Simple: just clip and continue
snake_x = np.clip(snake_x, 1, roi_w - 2)
snake_y = np.clip(snake_y, 1, roi_h - 2)
# Fit bounding box to refined contour
new_x1 = snake_x.min() + rx1
new_y1 = snake_y.min() + ry1
new_x2 = snake_x.max() + rx1
new_y2 = snake_y.max() + ry1
return np.array([new_x1, new_y1, new_x2, new_y2])
def compute_iou(b1, b2):
x1 = max(b1[0], b2[0]); y1 = max(b1[1], b2[1])
x2 = min(b1[2], b2[2]); y2 = min(b1[3], b2[3])
inter = max(0, x2-x1) * max(0, y2-y1)
a1 = (b1[2]-b1[0])*(b1[3]-b1[1])
a2 = (b2[2]-b2[0])*(b2[3]-b2[1])
return inter/(a1+a2-inter+1e-8)
def main():
yolo = YOLO('runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt')
val_img_dir = 'Data/Detection_dataset/images/val'
val_label_dir = 'Data/Detection_dataset/labels/val'
val_files = sorted([f for f in os.listdir(val_img_dir) if f.endswith('.jpg')])
# Test on first 5 images only (snake is slow)
test_files = val_files[:5]
print('Snake refinement test on 5 images:')
for img_file in test_files:
img_path = os.path.join(val_img_dir, img_file)
img = Image.open(img_path)
# YOLO predict
results = yolo.predict(img, imgsz=1536, conf=0.25, verbose=False)
if not results or len(results[0].boxes) == 0:
continue
boxes = results[0].boxes.xyxy.cpu().numpy()
# GT
lf = img_file.replace('.jpg', '.txt')
gt_boxes = []
with open(os.path.join(val_label_dir, lf)) as f:
for line in f:
p = line.strip().split()
if len(p) >= 5:
cx, cy, w, h = [float(x) for x in p[1:5]]
gt_boxes.append([
(cx-w/2)*img.size[0], (cy-h/2)*img.size[1],
(cx+w/2)*img.size[0], (cy+h/2)*img.size[1]
])
# Refine each box
for i, box in enumerate(boxes[:10]): # Max 10 per image
refined = snake_refine(img, box, n_iter=20)
if refined is not None:
# Measure improvement
best_gt_iou = 0
for gt in gt_boxes:
iou = compute_iou(box, gt)
if iou > best_gt_iou:
best_gt_iou = iou
best_ref_iou = 0
for gt in gt_boxes:
iou = compute_iou(refined, gt)
if iou > best_ref_iou:
best_ref_iou = iou
if best_ref_iou > best_gt_iou + 0.01:
delta = best_ref_iou - best_gt_iou
print(f' {img_file} box[{i}]: IoU {best_gt_iou:.3f} -> {best_ref_iou:.3f} (+{delta:.3f}) ✓')
elif best_ref_iou < best_gt_iou - 0.01:
delta = best_ref_iou - best_gt_iou
print(f' {img_file} box[{i}]: IoU {best_gt_iou:.3f} -> {best_ref_iou:.3f} ({delta:.3f}) ✗')
print('\nSnake test complete. CPU refinement works, no GPU needed.')
if __name__ == '__main__':
main()
|