| """Edge-Guided Bbox Refinement |
| ================================ |
| For each box edge, scan outward and find the image gradient peak. |
| More constrained than Snake, less likely to drift to wrong edges. |
| CPU-only, 0 training. |
| """ |
| import sys, os |
| import numpy as np |
| import cv2 |
| from PIL import Image |
| from tqdm import tqdm |
|
|
| PROJECT_DIR = '/home/user/goat' |
| os.chdir(PROJECT_DIR) |
| sys.path.insert(0, PROJECT_DIR) |
|
|
| from ultralytics import YOLO |
|
|
|
|
| def refine_edge(gray, x1, y1, x2, y2, edge, scan_range=10): |
| """Refine one edge of a bbox using gradient search. |
| |
| edge: 'left', 'right', 'top', 'bottom' |
| Returns: adjusted coordinate |
| """ |
| h, w = gray.shape |
| best_pos = None |
| best_grad = 0 |
|
|
| if edge == 'left': |
| x_center = int(x1) |
| for dx in range(-scan_range, scan_range + 1): |
| new_x = int(x1 + dx) |
| if new_x < 1 or new_x > w - 2: |
| continue |
| |
| strip = gray[max(0, int(y1)):min(h, int(y2)), new_x] |
| if len(strip) < 3: |
| continue |
| grad = np.abs(np.diff(strip)).mean() |
| if grad > best_grad: |
| best_grad = grad |
| best_pos = new_x |
| return best_pos if best_pos is not None else x1 |
|
|
| elif edge == 'right': |
| x_center = int(x2) |
| for dx in range(-scan_range, scan_range + 1): |
| new_x = int(x2 + dx) |
| if new_x < 1 or new_x > w - 2: |
| continue |
| strip = gray[max(0, int(y1)):min(h, int(y2)), new_x] |
| if len(strip) < 3: |
| continue |
| grad = np.abs(np.diff(strip)).mean() |
| if grad > best_grad: |
| best_grad = grad |
| best_pos = new_x |
| return best_pos if best_pos is not None else x2 |
|
|
| elif edge == 'top': |
| y_center = int(y1) |
| for dy in range(-scan_range, scan_range + 1): |
| new_y = int(y1 + dy) |
| if new_y < 1 or new_y > h - 2: |
| continue |
| strip = gray[new_y, max(0, int(x1)):min(w, int(x2))] |
| if len(strip) < 3: |
| continue |
| grad = np.abs(np.diff(strip)).mean() |
| if grad > best_grad: |
| best_grad = grad |
| best_pos = new_y |
| return best_pos if best_pos is not None else y1 |
|
|
| elif edge == 'bottom': |
| y_center = int(y2) |
| for dy in range(-scan_range, scan_range + 1): |
| new_y = int(y2 + dy) |
| if new_y < 1 or new_y > h - 2: |
| continue |
| strip = gray[new_y, max(0, int(x1)):min(w, int(x2))] |
| if len(strip) < 3: |
| continue |
| grad = np.abs(np.diff(strip)).mean() |
| if grad > best_grad: |
| best_grad = grad |
| best_pos = new_y |
| return best_pos if best_pos is not None else y2 |
|
|
| return x1 if edge in ('left', 'right') else y1 |
|
|
|
|
| def refine_box(img_gray, box, scan_range=8): |
| """Refine all 4 edges of a box using gradient search.""" |
| x1, y1, x2, y2 = box |
| w, h = x2 - x1, y2 - y1 |
| if w < 10 or h < 10: |
| return box |
|
|
| new_x1 = refine_edge(img_gray, x1, y1, x2, y2, 'left', scan_range) |
| new_x2 = refine_edge(img_gray, x1, y1, x2, y2, 'right', scan_range) |
| new_y1 = refine_edge(img_gray, x1, y1, x2, y2, 'top', scan_range) |
| new_y2 = refine_edge(img_gray, x1, y1, x2, y2, 'bottom', scan_range) |
|
|
| |
| min_w, min_h = w * 0.5, h * 0.5 |
| max_w, max_h = w * 1.5, h * 1.5 |
|
|
| if new_x2 - new_x1 < min_w: |
| mid = (new_x1 + new_x2) / 2 |
| new_x1 = mid - min_w / 2 |
| new_x2 = mid + min_w / 2 |
| if new_y2 - new_y1 < min_h: |
| mid = (new_y1 + new_y2) / 2 |
| new_y1 = mid - min_h / 2 |
| new_y2 = mid + min_h / 2 |
| if new_x2 - new_x1 > max_w: |
| mid = (new_x1 + new_x2) / 2 |
| new_x1 = mid - max_w / 2 |
| new_x2 = mid + max_w / 2 |
| if new_y2 - new_y1 > max_h: |
| mid = (new_y1 + new_y2) / 2 |
| new_y1 = mid - max_h / 2 |
| new_y2 = mid + max_h / 2 |
|
|
| return np.array([new_x1, new_y1, new_x2, new_y2]) |
|
|
|
|
| def compute_iou(b1, b2): |
| x1, y1 = max(b1[0], b2[0]), max(b1[1], b2[1]) |
| x2, y2 = min(b1[2], b2[2]), 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(): |
| model = 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_files = val_files[:5] |
| improved, degraded, total = 0, 0, 0 |
|
|
| print('Edge-Guided Refinement test:') |
| for img_file in test_files: |
| img_path = os.path.join(val_img_dir, img_file) |
| img = Image.open(img_path) |
| gray = np.array(img.convert('L'), dtype=float) |
|
|
| r = model.predict(img, imgsz=1536, conf=0.25, verbose=False) |
| if not r or len(r[0].boxes) == 0: |
| continue |
|
|
| boxes = r[0].boxes.xyxy.cpu().numpy() |
|
|
| |
| 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]]) |
|
|
| for box in boxes[:10]: |
| refined = refine_box(gray, box) |
| total += 1 |
|
|
| best_orig = max(compute_iou(box, gt) for gt in gt_boxes) if gt_boxes else 0 |
| best_ref = max(compute_iou(refined, gt) for gt in gt_boxes) if gt_boxes else 0 |
|
|
| if best_ref > best_orig + 0.005: |
| improved += 1 |
| elif best_ref < best_orig - 0.005: |
| degraded += 1 |
|
|
| print(f' Improved: {improved}/{total} ({improved/total*100:.0f}%)') |
| print(f' Degraded: {degraded}/{total} ({degraded/total*100:.0f}%)') |
| print(f' Unchanged: {total-improved-degraded}/{total}') |
|
|
| if improved > degraded: |
| print(' Edge refinement LOOKS PROMISING!') |
| else: |
| print(' Edge refinement not clearly beneficial') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
