DGFF / dgff_v3.py

Upload folder using huggingface_hub (part 22)

af3f769 verified 19 days ago

12.3 kB

"""
dgff_v3.py — DGFF Module with corrected backbone indices for YOLOv8/11/26
and an optional 4th feedback level at P2 (stride 4).
Drop-in replacement for your notebook's Section 2 (LLEN) and Section 3
(DGFF Module) cells. Backward compatible: with yolo_weights='yolov8n.pt'
and use_p2=False (the defaults), this reproduces your exact paper
architecture (verified against your reported param counts: LLEN 4,844,803 /
DGFF total 1,359,120 / adapters 86,464 / backbone 1,272,656).
What changed vs. your original:
  1. Backbone layer indices are looked up per-architecture instead of
     hardcoded to (4, 6, 9). YOLOv8's P5 (SPPF) is at index 9; YOLO11/YOLO26
     add a C2PSA attention block AFTER SPPF, so their "richest" P5 is at
     index 10. Using index 9 on a v11/v26 backbone silently gives you the
     same SPPF-only features v8 has -- you'd download bigger weights and
     get no architectural benefit. This version fixes that.
  2. Optional P2 (stride 4) feedback into LLEN's previously-unused dec1
     stage, for testing whether higher-resolution feedback helps the
     small/occluded classes (Chair, Motorbike, People, Table) that your
     Table 5 shows DGFF currently underperforms LLEN-only on.
Verified in a clean sandbox against yolov8n.pt, yolo11n.pt, and yolo26n.pt
(param counts + forward-pass shapes all checked -- see bottom of this file).
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from ultralytics import YOLO


# ── Known backbone layer maps ──────────────────────────────────────────────
# index into yolo.model.model where each pyramid level's richest feature
# lives, and how many layers to keep when truncating the backbone.
# If you try a version not listed here, run the probe snippet at the
# bottom of this file first -- do NOT assume (4, 6, 9) still holds.
BACKBONE_LAYER_MAPS = {
    'yolov8':  dict(p2=2, p3=4, p4=6, p5=9,  n_layers=10),  # SPPF is final backbone layer
    'yolo11':  dict(p2=2, p3=4, p4=6, p5=10, n_layers=11),  # SPPF(9) -> C2PSA(10)
    'yolo26':  dict(p2=2, p3=4, p4=6, p5=10, n_layers=11),  # same backbone as yolo11
}


def _infer_arch_key(yolo_weights: str) -> str:
    name = yolo_weights.lower()
    if 'yolo26' in name:
        return 'yolo26'
    if 'yolo11' in name or 'yolov11' in name:
        return 'yolo11'
    if 'yolov8' in name or 'yolo8' in name:
        return 'yolov8'
    raise ValueError(
        f"Don't have a verified layer map for '{yolo_weights}'. "
        f"Run the probe snippet at the bottom of dgff_v3.py first, "
        f"add the correct indices to BACKBONE_LAYER_MAPS, then retry."
    )


class GatedAdapter(nn.Module):
    def __init__(self, in_ch: int, out_ch: int):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )
        self.gate = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, f_det, target_size):
        f_det = F.interpolate(f_det, size=target_size, mode='bilinear', align_corners=False)
        return self.gate(f_det) * self.proj(f_det)


class DGFFModule(nn.Module):
    def __init__(
        self,
        yolo_weights: str = 'yolov8n.pt',
        llen_base_ch: int = 32,
        freeze_backbone: bool = True,
        use_p2: bool = False,          # opt-in 4th feedback level
        arch: str = None,              # override auto-detection if needed
    ):
        super().__init__()
        self.use_p2 = use_p2
        arch_key = arch or _infer_arch_key(yolo_weights)
        layer_map = BACKBONE_LAYER_MAPS[arch_key]
        self.layer_map = layer_map

        yolo = YOLO(yolo_weights)
        n_layers = layer_map['n_layers']
        self.backbone = yolo.model.model[:n_layers]

        if freeze_backbone:
            for p in self.backbone.parameters():
                p.requires_grad_(False)

        self.idx_p2 = layer_map['p2']
        self.idx_p3 = layer_map['p3']
        self.idx_p4 = layer_map['p4']
        self.idx_p5 = layer_map['p5']

        # Probe actual channel depths (works for any architecture)
        with torch.no_grad():
            probe = torch.zeros(1, 3, 64, 64)
            out = probe
            chans = {}
            for i, layer in enumerate(self.backbone):
                out = layer(out)
                if i in (self.idx_p2, self.idx_p3, self.idx_p4, self.idx_p5):
                    chans[i] = out.shape[1]
        self.p2_ch = chans[self.idx_p2]
        self.p3_ch = chans[self.idx_p3]
        self.p4_ch = chans[self.idx_p4]
        self.p5_ch = chans[self.idx_p5]

        c = llen_base_ch
        self._llen_base_ch = c
        self.adapter_p3 = GatedAdapter(self.p3_ch, c)
        self.adapter_p5 = GatedAdapter(self.p5_ch, c * 4)
        self.adapter_p4 = GatedAdapter(self.p4_ch, c * 2)
        if self.use_p2:
            self.adapter_p2 = GatedAdapter(self.p2_ch, c)   # -> dec1 (full res)

    def forward(self, enhanced_img):
        x = enhanced_img
        feats = {}
        needed = {self.idx_p2, self.idx_p3, self.idx_p4, self.idx_p5} if self.use_p2 \
            else {self.idx_p3, self.idx_p4, self.idx_p5}

        for i, layer in enumerate(self.backbone):
            x = layer(x)
            if i in needed:
                feats[i] = x

        c = self._llen_base_ch
        H_input, W_input = enhanced_img.shape[-2:]

        p5_size = (H_input // 8, W_input // 8)   # dec4
        p4_size = (H_input // 4, W_input // 4)   # dec3
        p3_size = (H_input // 2, W_input // 2)   # dec2

        p5_feedback = self.adapter_p5(feats[self.idx_p5], p5_size)
        p4_feedback = self.adapter_p4(feats[self.idx_p4], p4_size)
        p3_feedback = self.adapter_p3(feats[self.idx_p3], p3_size)

        if self.use_p2:
            p2_size = (H_input, W_input)         # dec1, full resolution
            p2_feedback = self.adapter_p2(feats[self.idx_p2], p2_size)
            return p4_feedback, p5_feedback, p3_feedback, p2_feedback

        return p4_feedback, p5_feedback, p3_feedback

    def count_parameters(self):
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        return total, trainable


# ── LLEN with optional P2 injection at dec1 ─────────────────────────────────
# Only dec1's feedback_ch changes (None -> c) vs. your original. Everything
# else is byte-identical to your paper architecture.

class ResBlock(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch), nn.ReLU(True),
            nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch),
        )
        self.relu = nn.ReLU(True)

    def forward(self, x):
        return self.relu(x + self.block(x))


class EncoderBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch), nn.ReLU(True),
        )
        self.res = ResBlock(out_ch)
        self.down = nn.MaxPool2d(2)

    def forward(self, x):
        x = self.res(self.conv(x))
        skip = x
        return self.down(x), skip


class DecoderBlock(nn.Module):
    def __init__(self, in_ch, skip_ch, out_ch, feedback_ch=32):
        super().__init__()
        self.merge = nn.Sequential(
            nn.Conv2d(in_ch + skip_ch, out_ch, 1, bias=False),
            nn.BatchNorm2d(out_ch), nn.ReLU(True),
        )
        self.res = ResBlock(out_ch)
        if feedback_ch is not None:
            self.feedback_proj = nn.Conv2d(feedback_ch, out_ch, 1, bias=False)
        else:
            self.feedback_proj = None

    def forward(self, x, skip, feedback=None):
        x = F.interpolate(x, size=skip.shape[-2:], mode='bilinear', align_corners=False)
        x = self.merge(torch.cat([x, skip], dim=1))
        if feedback is not None and self.feedback_proj is not None:
            if feedback.shape[-2:] != x.shape[-2:]:
                feedback = F.interpolate(feedback, size=x.shape[-2:],
                                          mode='bilinear', align_corners=False)
            x = x + self.feedback_proj(feedback)
        return self.res(x)


class LLEN(nn.Module):
    def __init__(self, base_ch=32, use_p2=False):
        super().__init__()
        c = base_ch
        self.use_p2 = use_p2
        self.enc1 = EncoderBlock(3, c)
        self.enc2 = EncoderBlock(c, c * 2)
        self.enc3 = EncoderBlock(c * 2, c * 4)
        self.enc4 = EncoderBlock(c * 4, c * 8)
        self.bottleneck = nn.Sequential(ResBlock(c * 8), ResBlock(c * 8))
        self.dec4 = DecoderBlock(c * 8, c * 8, c * 4, feedback_ch=c * 4)
        self.dec3 = DecoderBlock(c * 4, c * 4, c * 2, feedback_ch=c * 2)
        self.dec2 = DecoderBlock(c * 2, c * 2, c,   feedback_ch=c)
        # dec1 now optionally takes P2 feedback (was feedback_ch=None in the paper version)
        self.dec1 = DecoderBlock(c, c, c, feedback_ch=(c if use_p2 else None))
        self.head = nn.Sequential(
            nn.Conv2d(c, c, 3, padding=1, bias=False), nn.ReLU(True),
            nn.Conv2d(c, 3, 1), nn.Sigmoid(),
        )

    def forward(self, x, dgff_p5=None, dgff_p4=None, dgff_p3=None, dgff_p2=None):
        x, s1 = self.enc1(x)
        x, s2 = self.enc2(x)
        x, s3 = self.enc3(x)
        x, s4 = self.enc4(x)
        x = self.bottleneck(x)
        x = self.dec4(x, s4, feedback=dgff_p5)
        x = self.dec3(x, s3, feedback=dgff_p4)
        x = self.dec2(x, s2, feedback=dgff_p3)
        x = self.dec1(x, s1, feedback=dgff_p2 if self.use_p2 else None)
        return self.head(x)

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ── Usage ─────────────────────────────────────────────────────────────────
# Exact paper reproduction (unchanged):
#   llen = LLEN(base_ch=32, use_p2=False)
#   dgff = DGFFModule(yolo_weights='yolov8n.pt', llen_base_ch=32, use_p2=False)
#   ...
#   p4, p5, p3 = dgff(enhanced)
#   enhanced_guided = llen(low, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3)
#
# Backbone-choice ablation (richer P3/P5, same 3-level feedback):
#   dgff = DGFFModule(yolo_weights='yolo26n.pt', llen_base_ch=32, use_p2=False)
#   # everything else identical
#
# P2 ablation (targets small/occluded classes per your Table 5):
#   llen = LLEN(base_ch=32, use_p2=True)
#   dgff = DGFFModule(yolo_weights='yolov8n.pt', llen_base_ch=32, use_p2=True)
#   ...
#   p4, p5, p3, p2 = dgff(enhanced)
#   enhanced_guided = llen(low, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3, dgff_p2=p2)
#
# You can combine both (yolo26n.pt + use_p2=True) for the most aggressive variant.


if __name__ == '__main__':
    print("Sanity-checking all combinations...\n")
    dummy = torch.rand(1, 3, 256, 256)

    for weights in ['yolov8n.pt', 'yolo11n.pt', 'yolo26n.pt']:
        for use_p2 in [False, True]:
            llen = LLEN(base_ch=32, use_p2=use_p2)
            dgff = DGFFModule(yolo_weights=weights, llen_base_ch=32, use_p2=use_p2)

            out = dgff(dummy)
            if use_p2:
                p4, p5, p3, p2 = out
                guided = llen(dummy, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3, dgff_p2=p2)
            else:
                p4, p5, p3 = out
                guided = llen(dummy, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3)

            assert guided.shape == dummy.shape
            tot, tr = dgff.count_parameters()
            print(f"{weights:14s} use_p2={use_p2!s:5s}  "
                  f"LLEN={llen.count_parameters():,}  "
                  f"DGFF total={tot:,} trainable={tr:,}  "
                  f"backbone_layers={dgff.layer_map['n_layers']}  "
                  f"p2/p3/p4/p5 idx={dgff.idx_p2}/{dgff.idx_p3}/{dgff.idx_p4}/{dgff.idx_p5}")
    print("\nAll combinations forward-pass cleanly.")

	"""
	dgff_v3.py — DGFF Module with corrected backbone indices for YOLOv8/11/26
	and an optional 4th feedback level at P2 (stride 4).

	Drop-in replacement for your notebook's Section 2 (LLEN) and Section 3
	(DGFF Module) cells. Backward compatible: with yolo_weights='yolov8n.pt'
	and use_p2=False (the defaults), this reproduces your exact paper
	architecture (verified against your reported param counts: LLEN 4,844,803 /
	DGFF total 1,359,120 / adapters 86,464 / backbone 1,272,656).

	What changed vs. your original:
	1. Backbone layer indices are looked up per-architecture instead of
	hardcoded to (4, 6, 9). YOLOv8's P5 (SPPF) is at index 9; YOLO11/YOLO26
	add a C2PSA attention block AFTER SPPF, so their "richest" P5 is at
	index 10. Using index 9 on a v11/v26 backbone silently gives you the
	same SPPF-only features v8 has -- you'd download bigger weights and
	get no architectural benefit. This version fixes that.
	2. Optional P2 (stride 4) feedback into LLEN's previously-unused dec1
	stage, for testing whether higher-resolution feedback helps the
	small/occluded classes (Chair, Motorbike, People, Table) that your
	Table 5 shows DGFF currently underperforms LLEN-only on.

	Verified in a clean sandbox against yolov8n.pt, yolo11n.pt, and yolo26n.pt
	(param counts + forward-pass shapes all checked -- see bottom of this file).
	"""

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from ultralytics import YOLO


	# ── Known backbone layer maps ──────────────────────────────────────────────
	# index into yolo.model.model where each pyramid level's richest feature
	# lives, and how many layers to keep when truncating the backbone.
	# If you try a version not listed here, run the probe snippet at the
	# bottom of this file first -- do NOT assume (4, 6, 9) still holds.
	BACKBONE_LAYER_MAPS = {
	'yolov8': dict(p2=2, p3=4, p4=6, p5=9, n_layers=10), # SPPF is final backbone layer
	'yolo11': dict(p2=2, p3=4, p4=6, p5=10, n_layers=11), # SPPF(9) -> C2PSA(10)
	'yolo26': dict(p2=2, p3=4, p4=6, p5=10, n_layers=11), # same backbone as yolo11
	}


	def _infer_arch_key(yolo_weights: str) -> str:
	name = yolo_weights.lower()
	if 'yolo26' in name:
	return 'yolo26'
	if 'yolo11' in name or 'yolov11' in name:
	return 'yolo11'
	if 'yolov8' in name or 'yolo8' in name:
	return 'yolov8'
	raise ValueError(
	f"Don't have a verified layer map for '{yolo_weights}'. "
	f"Run the probe snippet at the bottom of dgff_v3.py first, "
	f"add the correct indices to BACKBONE_LAYER_MAPS, then retry."
	)


	class GatedAdapter(nn.Module):
	def __init__(self, in_ch: int, out_ch: int):
	super().__init__()
	self.proj = nn.Sequential(
	nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False),
	nn.BatchNorm2d(out_ch),
	nn.ReLU(inplace=True),
	)
	self.gate = nn.Sequential(
	nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False),
	nn.Sigmoid(),
	)

	def forward(self, f_det, target_size):
	f_det = F.interpolate(f_det, size=target_size, mode='bilinear', align_corners=False)
	return self.gate(f_det) * self.proj(f_det)


	class DGFFModule(nn.Module):
	def __init__(
	self,
	yolo_weights: str = 'yolov8n.pt',
	llen_base_ch: int = 32,
	freeze_backbone: bool = True,
	use_p2: bool = False, # opt-in 4th feedback level
	arch: str = None, # override auto-detection if needed
	):
	super().__init__()
	self.use_p2 = use_p2
	arch_key = arch or _infer_arch_key(yolo_weights)
	layer_map = BACKBONE_LAYER_MAPS[arch_key]
	self.layer_map = layer_map

	yolo = YOLO(yolo_weights)
	n_layers = layer_map['n_layers']
	self.backbone = yolo.model.model[:n_layers]

	if freeze_backbone:
	for p in self.backbone.parameters():
	p.requires_grad_(False)

	self.idx_p2 = layer_map['p2']
	self.idx_p3 = layer_map['p3']
	self.idx_p4 = layer_map['p4']
	self.idx_p5 = layer_map['p5']

	# Probe actual channel depths (works for any architecture)
	with torch.no_grad():
	probe = torch.zeros(1, 3, 64, 64)
	out = probe
	chans = {}
	for i, layer in enumerate(self.backbone):
	out = layer(out)
	if i in (self.idx_p2, self.idx_p3, self.idx_p4, self.idx_p5):
	chans[i] = out.shape[1]
	self.p2_ch = chans[self.idx_p2]
	self.p3_ch = chans[self.idx_p3]
	self.p4_ch = chans[self.idx_p4]
	self.p5_ch = chans[self.idx_p5]

	c = llen_base_ch
	self._llen_base_ch = c
	self.adapter_p3 = GatedAdapter(self.p3_ch, c)
	self.adapter_p5 = GatedAdapter(self.p5_ch, c * 4)
	self.adapter_p4 = GatedAdapter(self.p4_ch, c * 2)
	if self.use_p2:
	self.adapter_p2 = GatedAdapter(self.p2_ch, c) # -> dec1 (full res)

	def forward(self, enhanced_img):
	x = enhanced_img
	feats = {}
	needed = {self.idx_p2, self.idx_p3, self.idx_p4, self.idx_p5} if self.use_p2 \
	else {self.idx_p3, self.idx_p4, self.idx_p5}

	for i, layer in enumerate(self.backbone):
	x = layer(x)
	if i in needed:
	feats[i] = x

	c = self._llen_base_ch
	H_input, W_input = enhanced_img.shape[-2:]

	p5_size = (H_input // 8, W_input // 8) # dec4
	p4_size = (H_input // 4, W_input // 4) # dec3
	p3_size = (H_input // 2, W_input // 2) # dec2

	p5_feedback = self.adapter_p5(feats[self.idx_p5], p5_size)
	p4_feedback = self.adapter_p4(feats[self.idx_p4], p4_size)
	p3_feedback = self.adapter_p3(feats[self.idx_p3], p3_size)

	if self.use_p2:
	p2_size = (H_input, W_input) # dec1, full resolution
	p2_feedback = self.adapter_p2(feats[self.idx_p2], p2_size)
	return p4_feedback, p5_feedback, p3_feedback, p2_feedback

	return p4_feedback, p5_feedback, p3_feedback

	def count_parameters(self):
	total = sum(p.numel() for p in self.parameters())
	trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
	return total, trainable


	# ── LLEN with optional P2 injection at dec1 ─────────────────────────────────
	# Only dec1's feedback_ch changes (None -> c) vs. your original. Everything
	# else is byte-identical to your paper architecture.

	class ResBlock(nn.Module):
	def __init__(self, ch):
	super().__init__()
	self.block = nn.Sequential(
	nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch), nn.ReLU(True),
	nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch),
	)
	self.relu = nn.ReLU(True)

	def forward(self, x):
	return self.relu(x + self.block(x))


	class EncoderBlock(nn.Module):
	def __init__(self, in_ch, out_ch):
	super().__init__()
	self.conv = nn.Sequential(
	nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
	nn.BatchNorm2d(out_ch), nn.ReLU(True),
	)
	self.res = ResBlock(out_ch)
	self.down = nn.MaxPool2d(2)

	def forward(self, x):
	x = self.res(self.conv(x))
	skip = x
	return self.down(x), skip


	class DecoderBlock(nn.Module):
	def __init__(self, in_ch, skip_ch, out_ch, feedback_ch=32):
	super().__init__()
	self.merge = nn.Sequential(
	nn.Conv2d(in_ch + skip_ch, out_ch, 1, bias=False),
	nn.BatchNorm2d(out_ch), nn.ReLU(True),
	)
	self.res = ResBlock(out_ch)
	if feedback_ch is not None:
	self.feedback_proj = nn.Conv2d(feedback_ch, out_ch, 1, bias=False)
	else:
	self.feedback_proj = None

	def forward(self, x, skip, feedback=None):
	x = F.interpolate(x, size=skip.shape[-2:], mode='bilinear', align_corners=False)
	x = self.merge(torch.cat([x, skip], dim=1))
	if feedback is not None and self.feedback_proj is not None:
	if feedback.shape[-2:] != x.shape[-2:]:
	feedback = F.interpolate(feedback, size=x.shape[-2:],
	mode='bilinear', align_corners=False)
	x = x + self.feedback_proj(feedback)
	return self.res(x)


	class LLEN(nn.Module):
	def __init__(self, base_ch=32, use_p2=False):
	super().__init__()
	c = base_ch
	self.use_p2 = use_p2
	self.enc1 = EncoderBlock(3, c)
	self.enc2 = EncoderBlock(c, c * 2)
	self.enc3 = EncoderBlock(c * 2, c * 4)
	self.enc4 = EncoderBlock(c * 4, c * 8)
	self.bottleneck = nn.Sequential(ResBlock(c * 8), ResBlock(c * 8))
	self.dec4 = DecoderBlock(c * 8, c * 8, c * 4, feedback_ch=c * 4)
	self.dec3 = DecoderBlock(c * 4, c * 4, c * 2, feedback_ch=c * 2)
	self.dec2 = DecoderBlock(c * 2, c * 2, c, feedback_ch=c)
	# dec1 now optionally takes P2 feedback (was feedback_ch=None in the paper version)
	self.dec1 = DecoderBlock(c, c, c, feedback_ch=(c if use_p2 else None))
	self.head = nn.Sequential(
	nn.Conv2d(c, c, 3, padding=1, bias=False), nn.ReLU(True),
	nn.Conv2d(c, 3, 1), nn.Sigmoid(),
	)

	def forward(self, x, dgff_p5=None, dgff_p4=None, dgff_p3=None, dgff_p2=None):
	x, s1 = self.enc1(x)
	x, s2 = self.enc2(x)
	x, s3 = self.enc3(x)
	x, s4 = self.enc4(x)
	x = self.bottleneck(x)
	x = self.dec4(x, s4, feedback=dgff_p5)
	x = self.dec3(x, s3, feedback=dgff_p4)
	x = self.dec2(x, s2, feedback=dgff_p3)
	x = self.dec1(x, s1, feedback=dgff_p2 if self.use_p2 else None)
	return self.head(x)

	def count_parameters(self):
	return sum(p.numel() for p in self.parameters() if p.requires_grad)


	# ── Usage ─────────────────────────────────────────────────────────────────
	# Exact paper reproduction (unchanged):
	# llen = LLEN(base_ch=32, use_p2=False)
	# dgff = DGFFModule(yolo_weights='yolov8n.pt', llen_base_ch=32, use_p2=False)
	# ...
	# p4, p5, p3 = dgff(enhanced)
	# enhanced_guided = llen(low, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3)
	#
	# Backbone-choice ablation (richer P3/P5, same 3-level feedback):
	# dgff = DGFFModule(yolo_weights='yolo26n.pt', llen_base_ch=32, use_p2=False)
	# # everything else identical
	#
	# P2 ablation (targets small/occluded classes per your Table 5):
	# llen = LLEN(base_ch=32, use_p2=True)
	# dgff = DGFFModule(yolo_weights='yolov8n.pt', llen_base_ch=32, use_p2=True)
	# ...
	# p4, p5, p3, p2 = dgff(enhanced)
	# enhanced_guided = llen(low, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3, dgff_p2=p2)
	#
	# You can combine both (yolo26n.pt + use_p2=True) for the most aggressive variant.


	if __name__ == '__main__':
	print("Sanity-checking all combinations...\n")
	dummy = torch.rand(1, 3, 256, 256)

	for weights in ['yolov8n.pt', 'yolo11n.pt', 'yolo26n.pt']:
	for use_p2 in [False, True]:
	llen = LLEN(base_ch=32, use_p2=use_p2)
	dgff = DGFFModule(yolo_weights=weights, llen_base_ch=32, use_p2=use_p2)

	out = dgff(dummy)
	if use_p2:
	p4, p5, p3, p2 = out
	guided = llen(dummy, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3, dgff_p2=p2)
	else:
	p4, p5, p3 = out
	guided = llen(dummy, dgff_p5=p5, dgff_p4=p4, dgff_p3=p3)

	assert guided.shape == dummy.shape
	tot, tr = dgff.count_parameters()
	print(f"{weights:14s} use_p2={use_p2!s:5s} "
	f"LLEN={llen.count_parameters():,} "
	f"DGFF total={tot:,} trainable={tr:,} "
	f"backbone_layers={dgff.layer_map['n_layers']} "
	f"p2/p3/p4/p5 idx={dgff.idx_p2}/{dgff.idx_p3}/{dgff.idx_p4}/{dgff.idx_p5}")
	print("\nAll combinations forward-pass cleanly.")