File size: 7,923 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 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 | """BRN: Bbox Refinement Network
===============================
A lightweight CNN that takes a high-resolution ROI crop from the ORIGINAL IMAGE
and predicts sub-pixel (Ξcx, Ξcy, Ξw, Ξh) corrections to a bounding box.
Unlike YOLO which sees the world through stride-8 features, BRN sees RAW PIXELS
at 1:1 resolution. This gives it precision that YOLO physically cannot achieve.
Architecture: EfficientNet-B0-lite inspired, ~500K params.
Training: Generate infinite samples by adding noise to GT boxes.
Inference: YOLO detects β BRN refines β final boxes.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
import numpy as np
from PIL import Image
class BRN(nn.Module):
"""Lightweight Bbox Refinement Network.
Input: [B, 3, H, W] ROI crop (default 96x96)
Output: [B, 4] (Ξcx, Ξcy, Ξw, Ξh) in normalized ROI coordinates [-0.5, 0.5]
"""
def __init__(self, input_size=96):
super().__init__()
# Stem: 96 β 48
self.stem = nn.Sequential(
nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(32),
nn.SiLU(inplace=True),
)
# Block1: 48 β 24
self.block1 = nn.Sequential(
nn.Conv2d(32, 48, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(48),
nn.SiLU(inplace=True),
nn.Conv2d(48, 48, 3, padding=1, bias=False),
nn.BatchNorm2d(48),
nn.SiLU(inplace=True),
)
# Block2: 24 β 12
self.block2 = nn.Sequential(
nn.Conv2d(48, 64, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(64),
nn.SiLU(inplace=True),
nn.Conv2d(64, 64, 3, padding=1, bias=False),
nn.BatchNorm2d(64),
nn.SiLU(inplace=True),
)
# Block3: 12 β 6
self.block3 = nn.Sequential(
nn.Conv2d(64, 96, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(96),
nn.SiLU(inplace=True),
nn.Conv2d(96, 96, 3, padding=1, bias=False),
nn.BatchNorm2d(96),
nn.SiLU(inplace=True),
)
# Head
self.head = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(96, 64),
nn.SiLU(inplace=True),
nn.Dropout(0.1),
nn.Linear(64, 4),
)
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.zeros_(m.bias)
def forward(self, x):
x = self.stem(x)
x = self.block1(x)
x = self.block2(x)
x = self.block3(x)
x = self.head(x)
return x
class BRNRefiner:
"""BRN-based bounding box refiner.
Usage:
refiner = BRNRefiner('path/to/brn.pt', input_size=96)
refined_boxes = refiner.refine(image, yolo_boxes)
"""
def __init__(self, model_path, input_size=96, device='cuda'):
self.input_size = input_size
self.device = device
self.model = BRN(input_size=input_size).to(device)
if model_path:
self.model.load_state_dict(torch.load(model_path, map_location=device))
self.model.eval()
self.transform = T.Compose([
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
@torch.no_grad()
def refine(self, image, boxes):
"""Refine bounding boxes using BRN.
Args:
image: PIL Image or numpy array
boxes: [N, 4] in (x1, y1, x2, y2) pixel coordinates
Returns:
refined: [N, 4] refined boxes
"""
if len(boxes) == 0:
return boxes
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
iw, ih = image.size
refined = []
for box in boxes:
x1, y1, x2, y2 = box
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
w, h = max(x2 - x1, 1), max(y2 - y1, 1)
# Expand ROI by 50% to give context
roi_size = max(w, h) * 1.5
roi_size = min(roi_size, min(iw, ih)) # Don't exceed image
half = roi_size / 2
# ROI bounds
rx1 = max(0, cx - half)
ry1 = max(0, cy - half)
rx2 = min(iw, cx + half)
ry2 = min(ih, cy + half)
# Crop and resize
roi = image.crop((rx1, ry1, rx2, ry2))
roi_tensor = self.transform(roi.resize((self.input_size, self.input_size), Image.BICUBIC))
roi_tensor = roi_tensor.unsqueeze(0).to(self.device)
# Predict correction
delta = self.model(roi_tensor)[0].cpu().numpy() # [Ξcx, Ξcy, Ξw, Ξh]
# Decode: correction is in ROI-normalized space
roi_w, roi_h = rx2 - rx1, ry2 - ry1
dcx = delta[0] * roi_w
dcy = delta[1] * roi_h
dw = delta[2] * roi_w
dh = delta[3] * roi_h
new_cx = cx + dcx
new_cy = cy + dcy
new_w = max(1, w + dw)
new_h = max(1, h + dh)
new_x1 = max(0, new_cx - new_w / 2)
new_y1 = max(0, new_cy - new_h / 2)
new_x2 = min(iw, new_cx + new_w / 2)
new_y2 = min(ih, new_cy + new_h / 2)
refined.append([new_x1, new_y1, new_x2, new_y2])
return np.array(refined)
# ββ Training data generation ββββββββββββββββββββββββββββββββ
def generate_training_sample(image, gt_boxes, noise_level=0.05):
"""Generate a BRN training sample from one image.
For each GT box, add random noise and create a (noisy_box, GT_box) pair.
Args:
image: PIL Image
gt_boxes: [N, 4] in (cx, cy, w, h) normalized [0,1]
noise_level: std of Gaussian noise relative to box size
Returns:
rois: [M, 3, S, S] tensor of ROI crops
deltas: [M, 4] target corrections to recover GT from noisy box
"""
iw, ih = image.size
rois = []
deltas = []
for box in gt_boxes:
cx, cy, w, h = box
cx_px, cy_px = cx * iw, cy * ih
w_px, h_px = w * iw, h * ih
# Add Gaussian noise to all 4 params
std_c = noise_level * min(w_px, h_px) # Center noise proportional to size
std_s = noise_level * max(w_px, h_px) # Size noise
noisy_cx = cx_px + np.random.randn() * std_c
noisy_cy = cy_px + np.random.randn() * std_c
noisy_w = w_px * (1 + np.random.randn() * noise_level)
noisy_h = h_px * (1 + np.random.randn() * noise_level)
noisy_w = max(4, noisy_w)
noisy_h = max(4, noisy_h)
# Compute target delta: what correction needed from noisy β GT
delta_cx = (cx_px - noisy_cx) / max(w_px, h_px) # Normalize by box size
delta_cy = (cy_px - noisy_cy) / max(w_px, h_px)
delta_w = (w_px - noisy_w) / max(w_px, h_px)
delta_h = (h_px - noisy_h) / max(w_px, h_px)
# Expand ROI around noisy box
half = max(noisy_w, noisy_h) * 0.75
half = max(half, 16) # Minimum ROI size
rx1 = max(0, int(noisy_cx - half))
ry1 = max(0, int(noisy_cy - half))
rx2 = min(iw, int(noisy_cx + half))
ry2 = min(ih, int(noisy_cy + half))
if rx2 - rx1 < 8 or ry2 - ry1 < 8:
continue
try:
roi = image.crop((rx1, ry1, rx2, ry2))
rois.append(roi)
deltas.append([delta_cx, delta_cy, delta_w, delta_h])
except Exception:
continue
return rois, deltas
|