+
+```python
+from PIL import Image
+from zeromodels.models.efficientdet import (
+ EfficientDetDetect,
+ EfficientDetImageProcessor,
+)
+
+model = EfficientDetDetect.from_weights("zeromodels/efficientdet_d0")
+processor = EfficientDetImageProcessor.from_weights("zeromodels/efficientdet_d0")
+
+image = Image.open("assets/data/coco_bicycles.jpg").convert("RGB")
+inputs = processor(image)
+
+output = model(inputs["pixel_values"], training=False)
+# output["boxes"]: (1, 49104, 4)
+# output["scores"]: (1, 49104, 90)
+
+results = processor.post_process_object_detection(
+ output, threshold=0.3, target_sizes=inputs["original_sizes"]
+)[0]
+
+# Detections come back unordered, so sort by score for readability.
+detections = sorted(
+ zip(results["scores"], results["label_names"], results["boxes"]),
+ key=lambda d: -float(d[0]),
+)
+for score, name, box in detections:
+ print(f"{name:14s} {float(score):.3f} {[round(float(v)) for v in box]}")
+```
+
+```
+person 0.819 [214, 79, 345, 378]
+bicycle 0.782 [98, 208, 441, 412]
+person 0.746 [584, 137, 634, 239]
+person 0.743 [315, 106, 407, 364]
+person 0.710 [437, 126, 533, 318]
+person 0.674 [414, 128, 456, 220]
+motorcycle 0.582 [388, 187, 639, 350]
+motorcycle 0.493 [0, 238, 107, 384]
+motorcycle 0.485 [533, 162, 637, 251]
+potted plant 0.342 [209, 32, 295, 189]
+motorcycle 0.317 [532, 165, 596, 238]
+```
+
+Two riders on the bicycle, the bicycle itself, and the parked motorcycles down the street,
+all from a 3.9 M-parameter model. `threshold=0.3` is a reasonable default for D0; raise it to
+`0.4`–`0.5` to keep only the confident detections.
+
+### Batch Processing Multiple Images
+
+Pass a list of images and one `target_sizes` entry per image:
+
+
+
+```python
+from PIL import Image
+from zeromodels.models.efficientdet import (
+ EfficientDetDetect,
+ EfficientDetImageProcessor,
+)
+
+model = EfficientDetDetect.from_weights("zeromodels/efficientdet_d0")
+processor = EfficientDetImageProcessor.from_weights("zeromodels/efficientdet_d0")
+
+paths = ["assets/data/coco_cat_tv.jpg", "assets/data/coco_cats.jpg"]
+images = [Image.open(p).convert("RGB") for p in paths]
+
+inputs = processor(paths) # (2, 512, 512, 3)
+output = model(inputs["pixel_values"], training=False)
+
+results = processor.post_process_object_detection(
+ output, threshold=0.3, target_sizes=[(im.height, im.width) for im in images]
+)
+
+for path, result in zip(paths, results):
+ print(f"\n{path}")
+ detections = sorted(
+ zip(result["scores"], result["label_names"], result["boxes"]),
+ key=lambda d: -float(d[0]),
+ )
+ for score, name, box in detections:
+ print(f" {name:10s} {float(score):.3f} {[round(float(v)) for v in box]}")
+```
+
+```
+assets/data/coco_cat_tv.jpg
+ tv 0.779 [144, 9, 560, 288]
+ cat 0.749 [132, 187, 525, 348]
+ cup 0.676 [0, 420, 78, 478]
+ person 0.536 [201, 40, 352, 216]
+
+assets/data/coco_cats.jpg
+ cat 0.779 [345, 31, 630, 355]
+ cat 0.714 [16, 44, 318, 446]
+ remote 0.635 [40, 72, 173, 120]
+ bed 0.504 [15, 0, 640, 474]
+```
+
+Every image is letterboxed to the same square, so stacking is always safe; results come back
+as per-image lists, not a fixed-width tensor.
+
+## NMS Modes
+
+EfficientDet emits a score for **every** class at every anchor, so a single object can clear
+the threshold under more than one label (a dog also read as a cat, a truck also as a car).
+How the duplicates are resolved is the post-processor's job:
+
+- **`class_agnostic=True`** (the default, Google's `postprocess_global`): each anchor keeps
+ only its single highest-scoring class, then one NMS runs across everything. One object
+ yields one box. This is what the examples above use.
+- **`class_agnostic=False`** (`postprocess_per_class`): NMS runs independently per class, so
+ the same object can surface under several labels. Useful when you want the full per-class
+ ranking.
+
+```python
+results = processor.post_process_object_detection(
+ output, threshold=0.3, class_agnostic=False, target_sizes=inputs["original_sizes"]
+)[0]
+```
+
+Per-class NMS cannot suppress a dog box and a cat box over the same animal, because they are
+different classes; the class-agnostic default collapses them to the higher-scoring one.
+
+## Input Resolution
+
+EfficientDet is Functional, so the input shape is fixed when the model is constructed, and
+the anchor grid is generated for that size. Each variant has a native resolution (the table
+above), but the **weights are resolution-independent** — the conv, BatchNorm, and separable
+weights do not depend on the input size, and the only size-dependent tensor, the anchor grid,
+is a computed constant rather than a stored weight. So you can build any variant at any valid
+size and load the same checkpoint:
+
+```python
+model = EfficientDetDetect.from_weights("zeromodels/efficientdet_d0", image_size=768)
+processor = EfficientDetImageProcessor(image_size=768) # match the model
+```
+
+The `image_size=768` kwarg overrides the checkpoint's native 512; the 512-trained weights
+load unchanged, and the anchor count scales with the input:
+
+```
+512: 49104 anchors cat 0.78, cat 0.71, remote 0.63, bed 0.50
+640: 76725 anchors cat 0.83, cat 0.73, couch 0.64, remote 0.62
+896: 150381 anchors couch 0.78, cat 0.73, remote 0.73, remote 0.71
+```
+
+**The side must be divisible by 128.** The heads run on levels 3 through 7 (strides 8 to
+128), so the input has to survive seven halvings cleanly. Every native size (512, 640, 768,
+896, 1024, 1280, 1536) qualifies.
+
+**Accuracy is best near the trained resolution.** The model *runs* at any valid size, but
+scores and box quality drift as you move away from native (visible above at 640 and 896). For
+the best result, run each variant near its native size, or pick the variant whose native size
+matches your target.
+
+## Custom Class Names
+
+A model fine-tuned on your own dataset predicts your class indices, not COCO's. Pass the names
+through `label_names` so the result reads correctly:
+
+```python
+MY_CLASSES = ["cat", "dog", "bird"]
+
+results = processor.post_process_object_detection(
+ output,
+ threshold=0.3,
+ target_sizes=[(image.height, image.width)],
+ label_names=MY_CLASSES,
+)
+```
+
+Custom names are indexed directly by the 0-based class id (`label_names[c]`). The default
+COCO list is offset by one instead (`COCO_91_CLASSES[c + 1]`, since index 0 is the "N/A" id),
+so the post-processor handles the two cases automatically. Without `label_names` a custom
+model is silently mislabeled with COCO names; the integer `labels` are unaffected either way.
+
+## Data Format
+
+**Both the model and the processor support `channels_last` and `channels_first`.** Nothing is
+hard-coded to a layout, so the whole pipeline runs either way, and the same weights load into
+both — the conv, BatchNorm, and separable-conv weights are layout-independent. Detections are
+identical (boxes bit-exact, scores within backend rounding).
+
+They pick the format differently, which is the one thing to keep straight:
+
+| | How it picks the format |
+|---|---|
+| Processor | A `data_format` kwarg, per instance. `None` (the default) resolves to `keras.config.image_data_format()`. |
+| Model | Reads `keras.config.image_data_format()` when it is **constructed**. There is no `data_format` argument. |
+
+### Overriding the processor only
+
+```python
+EfficientDetImageProcessor(data_format="channels_last")("photo.jpg")["pixel_values"]
+# (1, 512, 512, 3)
+
+EfficientDetImageProcessor(data_format="channels_first")("photo.jpg")["pixel_values"]
+# (1, 3, 512, 512)
+```
+
+### Switching the whole pipeline
+
+Set the global format before constructing the model, and both sides agree:
+
+```python
+import keras
+
+keras.config.set_image_data_format("channels_first")
+
+model = EfficientDetDetect.from_weights("zeromodels/efficientdet_d0")
+processor = EfficientDetImageProcessor.from_weights("zeromodels/efficientdet_d0")
+
+inputs = processor(image)
+# inputs["pixel_values"] is (1, 3, 512, 512)
+output = model(inputs["pixel_values"], training=False)
+```
+
+Set it once at the top of a script, since already-built models keep the layout they were
+constructed with. The post-processor emits `xyxy` pixel boxes and class indices, which have no
+channel axis, so it is not format-sensitive.
\ No newline at end of file
diff --git a/docs/granite_speech5.md b/docs/granite_speech5.md
index 632c150f..5558acc5 100644
--- a/docs/granite_speech5.md
+++ b/docs/granite_speech5.md
@@ -22,12 +22,24 @@ repeats and dropping the blank token, turns the per-frame argmax into text.
### GraniteSpeech5CTC
```python
-GraniteSpeech5CTC(vocab_size=16384, hidden_size=1024, intermediate_size=4096,
- num_hidden_layers=16, num_attention_heads=8, head_dim=128,
- num_mel_bins=80, hidden_act="silu", max_position_embeddings=512,
- context_size=128, conv_kernel_size=7, conv_expansion_factor=2,
- subsample_layers=(0, 1), attention_bias=True, pad_token_id=0,
- name="GraniteSpeech5CTC")
+GraniteSpeech5CTC(
+ vocab_size=16384,
+ hidden_size=1024,
+ intermediate_size=4096,
+ num_hidden_layers=16,
+ num_attention_heads=8,
+ head_dim=128,
+ num_mel_bins=80,
+ hidden_act="silu",
+ max_position_embeddings=512,
+ context_size=128,
+ conv_kernel_size=7,
+ conv_expansion_factor=2,
+ subsample_layers=(0, 1),
+ attention_bias=True,
+ pad_token_id=0,
+ name="GraniteSpeech5CTC",
+)
```
The conformer encoder plus the CTC head (tied to the encoder's mid-layer self-conditioning
@@ -61,7 +73,7 @@ blank, and render text.
### GraniteSpeech5Model
```python
-GraniteSpeech5Model(...) # same encoder arguments as GraniteSpeech5CTC
+GraniteSpeech5Model(...) # same encoder arguments as GraniteSpeech5CTC
```
The conformer encoder backbone alone. Returns `last_hidden_state`
@@ -73,9 +85,16 @@ CTC projection, for feature extraction or a custom head.
### GraniteSpeech5FeatureExtractor
```python
-GraniteSpeech5FeatureExtractor(sampling_rate=16000, n_fft=512, win_length=400,
- hop_length=160, num_mel_bins=80, delta_win_length=3,
- logmel_floor_db=8.0, frame_stacking=2)
+GraniteSpeech5FeatureExtractor(
+ sampling_rate=16000,
+ n_fft=512,
+ win_length=400,
+ hop_length=160,
+ num_mel_bins=80,
+ delta_win_length=3,
+ logmel_floor_db=8.0,
+ frame_stacking=2,
+)
```
Pure-Keras log-mel(+delta) feature extractor. Computes a `torchaudio`-style mel
@@ -114,6 +133,7 @@ returns `input_features` + `attention_mask` (and padded CTC `labels` when `text`
```python
import os
+
os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow"
import soundfile as sf
@@ -125,9 +145,11 @@ from zeromodels.models.granite_speech5 import (
model = GraniteSpeech5CTC.from_weights("zeromodels/granite-speech-5.0-470m-turboctc")
features = GraniteSpeech5FeatureExtractor()
-tokenizer = GraniteSpeech5Tokenizer.from_weights("zeromodels/granite-speech-5.0-470m-turboctc")
+tokenizer = GraniteSpeech5Tokenizer.from_weights(
+ "zeromodels/granite-speech-5.0-470m-turboctc"
+)
-audio, sr = sf.read("speech.wav", dtype="float32") # 16 kHz mono
+audio, sr = sf.read("speech.wav", dtype="float32") # 16 kHz mono
inputs = features(audio, sampling_rate=sr)
predicted_ids = model.generate(inputs)
print(tokenizer.batch_decode(predicted_ids))
@@ -138,7 +160,9 @@ Or drive it end to end through the processor:
```python
from zeromodels.models.granite_speech5 import GraniteSpeech5Processor
-processor = GraniteSpeech5Processor.from_weights("zeromodels/granite-speech-5.0-470m-turboctc")
+processor = GraniteSpeech5Processor.from_weights(
+ "zeromodels/granite-speech-5.0-470m-turboctc"
+)
inputs = processor(audio=audio, sampling_rate=sr)
predicted_ids = model.generate(inputs)
print(processor.batch_decode(predicted_ids))
@@ -156,6 +180,10 @@ print(processor.batch_decode(predicted_ids))
Upstream and community safetensors load on the fly with the `hf:` prefix (no re-hosting):
```python
-model = GraniteSpeech5CTC.from_weights("hf:ibm-granite/granite-speech-5.0-470m-turboctc")
-tokenizer = GraniteSpeech5Tokenizer.from_hf("ibm-granite/granite-speech-5.0-470m-turboctc")
+model = GraniteSpeech5CTC.from_weights(
+ "hf:ibm-granite/granite-speech-5.0-470m-turboctc"
+)
+tokenizer = GraniteSpeech5Tokenizer.from_hf(
+ "ibm-granite/granite-speech-5.0-470m-turboctc"
+)
```
diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py
index 67b7cfed..d3e132bb 100644
--- a/tests/base/model_test_registry.py
+++ b/tests/base/model_test_registry.py
@@ -717,6 +717,21 @@
"pred_boxes": (2, 10, 4),
},
},
+ "EfficientDetDetect": {
+ "module": "zeromodels.models.efficientdet",
+ "model_cls": "EfficientDetDetect",
+ "model_type": "object_detection",
+ "init_kwargs": {
+ "backbone_name": "efficientnet_b0",
+ "image_size": 256,
+ "num_classes": 20,
+ },
+ "input_shape": (2, 256, 256, 3),
+ "expected_output_shape": {
+ "boxes": (2, 12276, 4),
+ "scores": (2, 12276, 20),
+ },
+ },
"RTDETRV2Detect": {
"module": "zeromodels.models.rt_detr_v2",
"model_cls": "RTDETRV2Detect",
@@ -4256,6 +4271,7 @@ def instantiate_model(config):
def create_test_input(config, batch_size=2, model=None):
from tests.fixtures import dummy_inputs
+
if config.get("multimodal_vlm") and model is not None:
return dummy_inputs.multimodal_vlm_input(model, batch_size=batch_size)
diff --git a/tests/fixtures/dummy_inputs.py b/tests/fixtures/dummy_inputs.py
index fc20d8ea..f7b098ca 100644
--- a/tests/fixtures/dummy_inputs.py
+++ b/tests/fixtures/dummy_inputs.py
@@ -309,13 +309,19 @@ def multimodal_vlm_input(model, batch_size=2, seq_len=6, vocab_size=128):
shape[-1] = ids.shape[1]
feed[n] = ops.convert_to_tensor(np.zeros(shape, dtype="int32"))
elif "video" in n and "pixel" in n:
- feed[n] = feed["pixel_values"] # mirror the image; an empty video breaks the ViT
+ feed[n] = feed[
+ "pixel_values"
+ ] # mirror the image; an empty video breaks the ViT
elif n == "video_grid_thw":
feed[n] = feed["image_grid_thw"]
elif "grid_thw" in n:
- feed[n] = ops.convert_to_tensor(np.tile(np.array(grid, "int32"), (batch_size, 1)))
+ feed[n] = ops.convert_to_tensor(
+ np.tile(np.array(grid, "int32"), (batch_size, 1))
+ )
elif "hws" in n:
- feed[n] = ops.convert_to_tensor(np.tile(np.array([gh, gw], "int32"), (batch_size, 1)))
+ feed[n] = ops.convert_to_tensor(
+ np.tile(np.array([gh, gw], "int32"), (batch_size, 1))
+ )
elif n == "image_sizes":
# Pixtral-style variable resolution: the true (H, W) of each image,
# which must equal the fed pixel_values' spatial dims (non-batch, non-3).
@@ -325,7 +331,9 @@ def multimodal_vlm_input(model, batch_size=2, seq_len=6, vocab_size=128):
if pv
else [spatial_default, spatial_default]
)
- feed[n] = ops.convert_to_tensor(np.tile(np.array(hw, "int32"), (batch_size, 1)))
+ feed[n] = ops.convert_to_tensor(
+ np.tile(np.array(hw, "int32"), (batch_size, 1))
+ )
elif "pixel" in n:
if len(sh) == 2:
feed[n] = ops.ones((batch_size * n_patches, sh[1]))
diff --git a/website/mkdocs.yml b/website/mkdocs.yml
index 323f9f7a..8bf1fd02 100644
--- a/website/mkdocs.yml
+++ b/website/mkdocs.yml
@@ -176,6 +176,7 @@ nav:
- DINO: dino.md
- DINOv2: dinov2.md
- DINOv3: dinov3.md
+ - EfficientDet: efficientdet.md
- EoMT: eomt.md
- Mask2Former: mask2former.md
- MaskFormer: maskformer.md
diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py
index 2a302536..a413d52c 100644
--- a/zeromodels/models/__init__.py
+++ b/zeromodels/models/__init__.py
@@ -23,6 +23,7 @@
dino,
dino_v2,
dino_v3,
+ efficientdet,
efficientformer,
efficientnet,
efficientnet_lite,
diff --git a/zeromodels/models/efficientdet/__init__.py b/zeromodels/models/efficientdet/__init__.py
new file mode 100644
index 00000000..febc340b
--- /dev/null
+++ b/zeromodels/models/efficientdet/__init__.py
@@ -0,0 +1,10 @@
+from .efficientdet_config import EfficientDetConfig
+from .efficientdet_image_processor import EfficientDetImageProcessor
+from .efficientdet_model import EfficientDetDetect, EfficientDetModel
+
+__all__ = [
+ "EfficientDetConfig",
+ "EfficientDetModel",
+ "EfficientDetDetect",
+ "EfficientDetImageProcessor",
+]
diff --git a/zeromodels/models/efficientdet/convert_efficientdet_h5_to_keras.py b/zeromodels/models/efficientdet/convert_efficientdet_h5_to_keras.py
new file mode 100644
index 00000000..cd12349b
--- /dev/null
+++ b/zeromodels/models/efficientdet/convert_efficientdet_h5_to_keras.py
@@ -0,0 +1,174 @@
+import re
+
+import h5py
+import numpy as np
+from keras import ops
+
+BACKBONE_PREFIX = "efficientnet-{bb}/efficientnet-{bb}"
+
+EFFICIENTDET_RECIPES = {
+ "efficientdet_d0": {
+ "backbone_name": "efficientnet_b0",
+ "image_size": 512,
+ "fpn_num_filters": 64,
+ "fpn_cell_repeats": 3,
+ "box_class_repeats": 3,
+ },
+ "efficientdet_d1": {
+ "backbone_name": "efficientnet_b1",
+ "image_size": 640,
+ "fpn_num_filters": 88,
+ "fpn_cell_repeats": 4,
+ "box_class_repeats": 3,
+ },
+ "efficientdet_d2": {
+ "backbone_name": "efficientnet_b2",
+ "image_size": 768,
+ "fpn_num_filters": 112,
+ "fpn_cell_repeats": 5,
+ "box_class_repeats": 3,
+ },
+ "efficientdet_d3": {
+ "backbone_name": "efficientnet_b3",
+ "image_size": 896,
+ "fpn_num_filters": 160,
+ "fpn_cell_repeats": 6,
+ "box_class_repeats": 4,
+ },
+ "efficientdet_d4": {
+ "backbone_name": "efficientnet_b4",
+ "image_size": 1024,
+ "fpn_num_filters": 224,
+ "fpn_cell_repeats": 7,
+ "box_class_repeats": 4,
+ },
+ "efficientdet_d5": {
+ "backbone_name": "efficientnet_b5",
+ "image_size": 1280,
+ "fpn_num_filters": 288,
+ "fpn_cell_repeats": 7,
+ "box_class_repeats": 4,
+ },
+ "efficientdet_d6": {
+ "backbone_name": "efficientnet_b6",
+ "image_size": 1280,
+ "fpn_num_filters": 384,
+ "fpn_cell_repeats": 8,
+ "box_class_repeats": 5,
+ "fpn_weight_method": "sum",
+ },
+ "efficientdet_d7": {
+ "backbone_name": "efficientnet_b6",
+ "image_size": 1536,
+ "fpn_num_filters": 384,
+ "fpn_cell_repeats": 8,
+ "box_class_repeats": 5,
+ "anchor_scale": 5.0,
+ "fpn_weight_method": "sum",
+ },
+}
+
+EFFICIENTDET_H5_URLS = {
+ name: (
+ "https://storage.googleapis.com/cloud-tpu-checkpoints/efficientdet/"
+ f"coco2/{name.replace('_', '-')}.h5"
+ )
+ for name in EFFICIENTDET_RECIPES
+}
+
+
+def block_order(paths):
+ seen = []
+ for p in paths:
+ mo = re.match(r"(blocks_\d+_\d+)", p)
+ if mo and mo.group(1) not in seen:
+ seen.append(mo.group(1))
+ seen.sort(key=lambda s: tuple(int(x) for x in s.split("_")[1:]))
+ return {name: i for i, name in enumerate(seen)}
+
+
+def map_backbone_block(path, bb_prefix, block_n, has_expand):
+ rest = re.sub(r"^blocks_\d+_\d+_", "", path)
+ sub, var = rest.split("/", 1)
+ if sub == "conv2d_1": # expand conv
+ tgt = "conv2d"
+ elif sub == "conv2d_2": # project conv
+ tgt = "conv2d_1" if has_expand else "conv2d"
+ elif sub == "dwconv2d":
+ return f"{bb_prefix}/blocks_{block_n}/depthwise_conv2d/depthwise_kernel"
+ elif sub == "se_conv_reduce":
+ tgt = "se/conv2d"
+ elif sub == "se_conv_expand":
+ tgt = "se/conv2d_1"
+ elif sub == "batchnorm_1": # expand BN
+ tgt = "tpu_batch_normalization"
+ elif sub == "batchnorm_2": # depthwise BN
+ tgt = "tpu_batch_normalization_1" if has_expand else "tpu_batch_normalization"
+ elif sub == "batchnorm_3": # project BN
+ tgt = "tpu_batch_normalization_2" if has_expand else "tpu_batch_normalization_1"
+ else:
+ raise ValueError(f"unmapped block sublayer {sub!r} in {path}")
+ return f"{bb_prefix}/blocks_{block_n}/{tgt}/{var}"
+
+
+def keras_path_to_h5(path, bb_prefix, order, expand_flags, num_levels):
+ seg0 = path.split("/")[0]
+ if seg0 == "decode_boxes": # anchor constant, not in the checkpoint
+ return None
+ if path.startswith("conv_stem/"):
+ return f"{bb_prefix}/stem/conv2d/{path.split('/', 1)[1]}"
+ if path.startswith("batchnorm_1/"):
+ return f"{bb_prefix}/stem/tpu_batch_normalization/{path.split('/', 1)[1]}"
+ mo = re.match(r"(blocks_\d+_\d+)", path)
+ if mo:
+ block = mo.group(1)
+ return map_backbone_block(path, bb_prefix, order[block], expand_flags[block])
+ if seg0.startswith("resample_p"):
+ return f"{seg0}/{path}"
+ if seg0 in ("class_net", "box_net"):
+ return f"{seg0}/{path}"
+ if seg0 == "fpn_cells":
+ m2 = re.match(r"fpn_cells/cell_(\d+)_fnode(\d+)/(.+)", path)
+ r, i, rest = m2.group(1), int(m2.group(2)), m2.group(3)
+ n = num_levels + i
+ base = f"fpn_cells/fpn_cells/cell_{r}/fnode{i}"
+ if rest.startswith("WSM"):
+ return f"{base}/{rest}"
+ if rest.startswith("conv/") or rest.startswith("bn/"):
+ return f"{base}/op_after_combine{n}/{rest}"
+ mr = re.match(r"resample_(\d+)_(\d+)/(.+)", rest)
+ if mr:
+ return f"{base}/resample_{mr.group(1)}_{mr.group(2)}_{n}/{mr.group(3)}"
+ raise ValueError(f"unmapped fpn sublayer {rest!r} in {path}")
+ raise ValueError(f"unmapped path {path!r}")
+
+
+def transfer_efficientdet_weights(keras_model, h5_path, backbone=None):
+ if backbone is None:
+ backbone = keras_model.backbone_name.replace("_", "-")
+ paths = [w.path for w in keras_model.weights]
+ order = block_order(paths)
+ expand_flags = {
+ b: any(p.startswith(f"{b}_conv2d_1/") for p in paths) for b in order
+ }
+ num_levels = keras_model.max_level - keras_model.min_level + 1
+ bb_prefix = BACKBONE_PREFIX.format(bb=backbone.replace("efficientnet-", ""))
+
+ transferred = 0
+ with h5py.File(h5_path, "r") as f:
+ for w in keras_model.weights:
+ key = keras_path_to_h5(w.path, bb_prefix, order, expand_flags, num_levels)
+ if key is None:
+ continue
+ dataset = key + ":0"
+ if dataset not in f:
+ raise KeyError(f"{w.path} -> {dataset} not in {h5_path}")
+ value = np.asarray(f[dataset])
+ if tuple(value.shape) != tuple(w.shape):
+ raise ValueError(
+ f"shape mismatch {w.path}: model {tuple(w.shape)} "
+ f"vs h5 {tuple(value.shape)}"
+ )
+ w.assign(ops.convert_to_tensor(value))
+ transferred += 1
+ return transferred
diff --git a/zeromodels/models/efficientdet/efficientdet_config.py b/zeromodels/models/efficientdet/efficientdet_config.py
new file mode 100644
index 00000000..b31c2e78
--- /dev/null
+++ b/zeromodels/models/efficientdet/efficientdet_config.py
@@ -0,0 +1,80 @@
+from zeromodels.base import BaseConfig
+
+
+class EfficientDetConfig(BaseConfig):
+ r"""Configuration for EfficientDet (Google AutoML) object detection.
+
+ EfficientDet = an EfficientNet backbone -> a weighted bidirectional feature
+ pyramid (BiFPN, repeated ``fpn_cell_repeats`` times) -> shared class + box
+ prediction heads run over pyramid levels ``min_level``..``max_level``. Defaults
+ describe **EfficientDet-D0**; the D1-D7 variants override a handful of fields.
+ Pretrained variants load their architecture from the hosted ``zm_config.json`` via
+ ``from_weights``; the per-variant overrides used only when converting Google's
+ config-less ``.h5`` live in ``convert_efficientdet_h5_to_keras.EFFICIENTDET_RECIPES``.
+
+ Args:
+ backbone_name (`str`, *optional*, defaults to `"efficientnet_b0"`):
+ zeromodels EfficientNet backbone variant (the ``tf_efficientnet`` flavor
+ matches Google's backbone). D0->b0, D1->b1, ... D6/D7->b6.
+ image_size (`int`, *optional*, defaults to `512`):
+ Square input resolution.
+ num_classes (`int`, *optional*, defaults to `90`):
+ COCO class count (background excluded).
+ min_level / max_level (`int`, *optional*, defaults to `3` / `7`):
+ Pyramid levels the BiFPN + heads run over.
+ num_scales (`int`, *optional*, defaults to `3`):
+ Anchor scales per level; ``num_anchors = num_scales * len(aspect_ratios)``.
+ aspect_ratios (`tuple`, *optional*, defaults to `(1.0, 2.0, 0.5)`):
+ Anchor aspect ratios.
+ anchor_scale (`float`, *optional*, defaults to `4.0`):
+ Base anchor size in units of the feature stride.
+ fpn_num_filters (`int`, *optional*, defaults to `64`):
+ BiFPN (and head) channel width.
+ fpn_cell_repeats (`int`, *optional*, defaults to `3`):
+ Number of stacked BiFPN cells.
+ box_class_repeats (`int`, *optional*, defaults to `3`):
+ Shared conv layers in each of the class / box heads.
+ act_type (`str`, *optional*, defaults to `"swish"`):
+ Activation used throughout the BiFPN and heads.
+ separable_conv (`bool`, *optional*, defaults to `True`):
+ Use depthwise-separable 3x3 convs (as EfficientDet does).
+ apply_bn_for_resampling (`bool`, *optional*, defaults to `True`):
+ BatchNorm after the 1x1 channel-match in resampling.
+ conv_after_downsample (`bool`, *optional*, defaults to `False`):
+ Apply the 1x1 channel-match after (vs before) the downsample pool.
+ conv_bn_act_pattern (`bool`, *optional*, defaults to `False`):
+ `False` => act -> conv -> BN in a BiFPN node; `True` => conv -> BN -> act.
+ fpn_weight_method (`str`, *optional*, defaults to `"fastattn"`):
+ BiFPN edge fusion: `"fastattn"` (fast normalized), `"attn"` (softmax),
+ or `"sum"`.
+ survival_prob (`float`, *optional*):
+ Drop-connect keep prob in the heads (inference ignores it).
+
+ Example:
+
+ ```python
+ >>> from zeromodels.models.efficientdet import EfficientDetConfig
+
+ >>> config = EfficientDetConfig(backbone_name="efficientnet_b2", image_size=768)
+ ```"""
+
+ model_type = "efficientdet"
+
+ backbone_name: str = "efficientnet_b0"
+ image_size: int = 512
+ num_classes: int = 90
+ min_level: int = 3
+ max_level: int = 7
+ num_scales: int = 3
+ aspect_ratios: tuple = (1.0, 2.0, 0.5)
+ anchor_scale: float = 4.0
+ fpn_num_filters: int = 64
+ fpn_cell_repeats: int = 3
+ box_class_repeats: int = 3
+ act_type: str = "swish"
+ separable_conv: bool = True
+ apply_bn_for_resampling: bool = True
+ conv_after_downsample: bool = False
+ conv_bn_act_pattern: bool = False
+ fpn_weight_method: str = "fastattn"
+ survival_prob: float = None
diff --git a/zeromodels/models/efficientdet/efficientdet_image_processor.py b/zeromodels/models/efficientdet/efficientdet_image_processor.py
new file mode 100644
index 00000000..5dc132d6
--- /dev/null
+++ b/zeromodels/models/efficientdet/efficientdet_image_processor.py
@@ -0,0 +1,215 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+import keras
+import numpy as np
+from keras import ops
+from PIL import Image
+
+from zeromodels.base import BaseImageProcessor
+from zeromodels.utils.image_util import get_data_format, load_image
+from zeromodels.utils.labels_util import COCO_91_CLASSES
+
+from .efficientdet_layers import EfficientDetNMS
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class EfficientDetImageProcessor(BaseImageProcessor):
+ """Preprocess images for EfficientDet inference.
+
+ EfficientDet uses an aspect-preserving letterbox: the image is scaled by
+ ``image_size / max(height, width)``, ImageNet-normalized, then zero-padded on the
+ bottom / right to a square ``image_size x image_size``. The scale factor is
+ returned so detections can be mapped back to the original image (also recomputed
+ from ``target_sizes`` in
+ :meth:`EfficientDetDetect.post_process_object_detection`).
+
+ Args:
+ image_size: Square target size. Defaults to ``512`` (EfficientDet-D0).
+ resample: Interpolation method. Defaults to ``"bilinear"``.
+ image_mean: Per-channel mean. Defaults to ImageNet mean.
+ image_std: Per-channel std. Defaults to ImageNet std.
+ rescale_factor: Scale applied before normalization. Defaults to ``1/255``.
+ return_tensor: If True return a Keras tensor, else a numpy array.
+ data_format: ``"channels_first"`` / ``"channels_last"``; ``None`` resolves to
+ the Keras default.
+ """
+
+ def __init__(
+ self,
+ image_size: int = 512,
+ resample: str = "bilinear",
+ image_mean: Optional[Tuple[float, ...]] = None,
+ image_std: Optional[Tuple[float, ...]] = None,
+ rescale_factor: float = 1 / 255,
+ return_tensor: bool = True,
+ data_format: Optional[str] = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.image_size = image_size
+ self.resample = resample
+ self.image_mean = (
+ image_mean if image_mean is not None else (0.485, 0.456, 0.406)
+ )
+ self.image_std = image_std if image_std is not None else (0.229, 0.224, 0.225)
+ self.rescale_factor = rescale_factor
+ self.return_tensor = return_tensor
+ self.data_format = data_format
+
+ def __call__(self, image):
+ return self.call(image)
+
+ def preprocess_one(self, image):
+ arr = load_image(image)
+ orig_h, orig_w = int(arr.shape[0]), int(arr.shape[1])
+ scale = self.image_size / float(max(orig_h, orig_w))
+ new_h = int(orig_h * scale)
+ new_w = int(orig_w * scale)
+
+ # Preprocess entirely in channels_last (NHWC), then transpose at the end, so
+ # the base resize/normalize helpers are not misled by a channels_first global.
+ t = ops.convert_to_tensor(arr, dtype="float32")
+ t = ops.expand_dims(t, axis=0)
+ t = self.resize(
+ t,
+ (new_h, new_w),
+ interpolation=self.resample,
+ antialias=True,
+ data_format="channels_last",
+ )
+ t = self.rescale_and_normalize(
+ t,
+ do_rescale=True,
+ scale=self.rescale_factor,
+ do_normalize=True,
+ mean=self.image_mean,
+ std=self.image_std,
+ data_format="channels_last",
+ )
+ pad_h = self.image_size - new_h
+ pad_w = self.image_size - new_w
+ t = ops.pad(t, [(0, 0), (0, pad_h), (0, pad_w), (0, 0)], constant_values=0.0)
+ # data_format=None follows keras.config.image_data_format(), so the output
+ # layout matches a model built under the same global setting.
+ if get_data_format(self.data_format) == "channels_first":
+ t = ops.transpose(t, (0, 3, 1, 2))
+ return t, scale, (orig_h, orig_w)
+
+ def call(
+ self, image: Union[str, np.ndarray, Image.Image, List]
+ ) -> Dict[str, Union[keras.KerasTensor, np.ndarray]]:
+ items = list(image) if isinstance(image, (list, tuple)) else [image]
+ tensors, scales, sizes = [], [], []
+ for item in items:
+ t, scale, size = self.preprocess_one(item)
+ tensors.append(t)
+ scales.append(scale)
+ sizes.append(size)
+ pixel_values = ops.concatenate(tensors, axis=0)
+ if not self.return_tensor:
+ pixel_values = ops.convert_to_numpy(pixel_values)
+ return {
+ "pixel_values": pixel_values,
+ "scales": scales,
+ "original_sizes": sizes,
+ }
+
+ def post_process_object_detection(
+ self,
+ outputs,
+ threshold: float = 0.3,
+ iou_threshold: float = 0.5,
+ max_detections: int = 100,
+ class_agnostic: bool = True,
+ target_sizes=None,
+ label_names=None,
+ ):
+ """Turn raw :class:`EfficientDetDetect` outputs into per-image detections.
+
+ Applies per-class hard NMS (via :class:`EfficientDetNMS`) to the decoded
+ boxes and sigmoid scores, then, when ``target_sizes`` is given, undoes the
+ letterbox (divides by ``image_size / max(orig_h, orig_w)``) and clips boxes
+ to each original image.
+
+ Args:
+ outputs: The detector output dict ``{"boxes": (B, N, 4),
+ "scores": (B, N, num_classes)}``. Boxes are ``[ymin, xmin, ymax,
+ xmax]`` in ``image_size`` pixel coordinates.
+ threshold: Minimum score to keep a detection.
+ iou_threshold: NMS IoU threshold.
+ max_detections: Maximum detections kept per image.
+ class_agnostic: If True (default, Google's ``postprocess_global``) run one
+ NMS across all classes so an object yields a single box; if False, run
+ NMS per class.
+ target_sizes: List of ``(height, width)`` original image sizes, one per
+ batch item, used to rescale boxes back to the original image. If
+ ``None``, boxes stay in ``image_size`` coordinates.
+ label_names: Class-name list indexed by detector class id. Defaults to
+ the 90 COCO categories (class id ``c`` -> ``COCO_91_CLASSES[c + 1]``).
+
+ Returns:
+ A list (one dict per image) with ``"scores"``, ``"labels"``,
+ ``"label_names"`` and ``"boxes"`` (``[xmin, ymin, xmax, ymax]``).
+ """
+ nms = EfficientDetNMS(
+ iou_threshold=iou_threshold,
+ score_threshold=threshold,
+ max_detections=max_detections,
+ class_agnostic=class_agnostic,
+ )
+ detections = nms(outputs["boxes"], outputs["scores"])
+ if label_names is None:
+ names, offset = COCO_91_CLASSES, 1
+ else:
+ names, offset = label_names, 0
+
+ results = []
+ for i in range(int(detections.shape[0])):
+ det = detections[i]
+ num_valid = int(ops.sum(ops.cast(det[:, 4] > 0.0, "int32")))
+ det = det[:num_valid]
+ ymin, xmin, ymax, xmax = det[:, 0], det[:, 1], det[:, 2], det[:, 3]
+ scores = det[:, 4]
+ labels = ops.cast(det[:, 5], "int32")
+ if target_sizes is not None:
+ orig_h, orig_w = target_sizes[i]
+ scale = self.image_size / float(max(orig_h, orig_w))
+ ymin, xmin, ymax, xmax = (
+ ymin / scale,
+ xmin / scale,
+ ymax / scale,
+ xmax / scale,
+ )
+ ymin = ops.clip(ymin, 0, orig_h)
+ ymax = ops.clip(ymax, 0, orig_h)
+ xmin = ops.clip(xmin, 0, orig_w)
+ xmax = ops.clip(xmax, 0, orig_w)
+ boxes_xyxy = ops.stack([xmin, ymin, xmax, ymax], axis=-1)
+ mapped = [
+ names[c + offset] if (c + offset) < len(names) else f"class_{c}"
+ for c in (int(x) for x in labels)
+ ]
+ results.append(
+ {
+ "scores": ops.convert_to_numpy(scores),
+ "labels": ops.convert_to_numpy(labels),
+ "label_names": mapped,
+ "boxes": ops.convert_to_numpy(boxes_xyxy),
+ }
+ )
+ return results
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "image_size": self.image_size,
+ "resample": self.resample,
+ "image_mean": self.image_mean,
+ "image_std": self.image_std,
+ "rescale_factor": self.rescale_factor,
+ "return_tensor": self.return_tensor,
+ "data_format": self.data_format,
+ }
+ )
+ return config
diff --git a/zeromodels/models/efficientdet/efficientdet_layers.py b/zeromodels/models/efficientdet/efficientdet_layers.py
new file mode 100644
index 00000000..33244f30
--- /dev/null
+++ b/zeromodels/models/efficientdet/efficientdet_layers.py
@@ -0,0 +1,741 @@
+import math
+
+import keras
+from keras import layers, ops
+
+from zeromodels.utils.image_util import get_data_format
+
+
+def channel_axis(data_format):
+ return 1 if data_format == "channels_first" else -1
+
+
+def spatial_dims(shape, data_format):
+ """(height, width) of a 4D feature-map shape under ``data_format``."""
+ if data_format == "channels_first":
+ return int(shape[2]), int(shape[3])
+ return int(shape[1]), int(shape[2])
+
+
+def act(x, act_type="swish"):
+ if act_type == "swish":
+ return ops.silu(x)
+ if act_type == "relu":
+ return ops.relu(x)
+ return layers.Activation(act_type)(x)
+
+
+def sep_conv(filters, kernel_size=3, use_bias=True, data_format=None, name=None):
+ return layers.SeparableConv2D(
+ filters,
+ kernel_size=kernel_size,
+ padding="same",
+ depth_multiplier=1,
+ use_bias=use_bias,
+ data_format=data_format,
+ name=name,
+ )
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class EfficientDetResample(layers.Layer):
+ """Resample a feature map to a target level: 1x1 conv (only when the channel
+ count changes) + BatchNorm, then max-pool (downsample) or nearest-neighbour
+ upsample to the target spatial size. ``conv_after_downsample`` moves the 1x1
+ to after the pool."""
+
+ def __init__(
+ self,
+ target_channels,
+ apply_bn=True,
+ conv_after_downsample=False,
+ data_format=None,
+ name=None,
+ ):
+ super().__init__(name=name)
+ self.target_channels = target_channels
+ self.apply_bn = apply_bn
+ self.conv_after_downsample = conv_after_downsample
+ self.data_format = get_data_format(data_format)
+ self.conv2d = None
+ self.bn = None
+
+ def build(self, input_shape):
+ in_channels = input_shape[channel_axis(self.data_format)]
+ if in_channels != self.target_channels:
+ self.conv2d = layers.Conv2D(
+ self.target_channels,
+ 1,
+ padding="same",
+ data_format=self.data_format,
+ name="conv2d",
+ )
+ if self.apply_bn:
+ self.bn = layers.BatchNormalization(
+ axis=channel_axis(self.data_format), name="bn"
+ )
+ self.built = True
+
+ def apply_1x1(self, feat, training):
+ if self.conv2d is not None:
+ feat = self.conv2d(feat)
+ if self.bn is not None:
+ feat = self.bn(feat, training=training)
+ return feat
+
+ def call(self, feat, target_height, target_width, training=False):
+ height, width = spatial_dims(feat.shape, self.data_format)
+ if height > target_height and width > target_width:
+ if not self.conv_after_downsample:
+ feat = self.apply_1x1(feat, training)
+ h_stride = (height - 1) // target_height + 1
+ w_stride = (width - 1) // target_width + 1
+ feat = layers.MaxPooling2D(
+ pool_size=(h_stride + 1, w_stride + 1),
+ strides=(h_stride, w_stride),
+ padding="same",
+ data_format=self.data_format,
+ )(feat)
+ if self.conv_after_downsample:
+ feat = self.apply_1x1(feat, training)
+ else:
+ feat = self.apply_1x1(feat, training)
+ if height < target_height or width < target_width:
+ feat = ops.image.resize(
+ feat,
+ (target_height, target_width),
+ interpolation="nearest",
+ data_format=self.data_format,
+ )
+ return feat
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "target_channels": self.target_channels,
+ "apply_bn": self.apply_bn,
+ "conv_after_downsample": self.conv_after_downsample,
+ "data_format": self.data_format,
+ }
+ )
+ return config
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class FNode(layers.Layer):
+ """One BiFPN node: resample each input feature to this node's level, fuse them
+ with a weighted sum (``fastattn`` = normalized ReLU weights, ``attn`` = softmax,
+ ``sum`` = plain), then act -> separable 3x3 conv -> BN. Appends the new feature
+ to the running list."""
+
+ def __init__(
+ self,
+ feat_level,
+ inputs_offsets,
+ fpn_num_filters,
+ weight_method,
+ act_type,
+ apply_bn_for_resampling,
+ conv_after_downsample,
+ conv_bn_act_pattern,
+ data_format=None,
+ name=None,
+ ):
+ super().__init__(name=name)
+ self.feat_level = feat_level
+ self.inputs_offsets = list(inputs_offsets)
+ self.fpn_num_filters = fpn_num_filters
+ self.weight_method = weight_method
+ self.act_type = act_type
+ self.apply_bn_for_resampling = apply_bn_for_resampling
+ self.conv_after_downsample = conv_after_downsample
+ self.conv_bn_act_pattern = conv_bn_act_pattern
+ self.data_format = get_data_format(data_format)
+ self.resamples = [
+ EfficientDetResample(
+ fpn_num_filters,
+ apply_bn=apply_bn_for_resampling,
+ conv_after_downsample=conv_after_downsample,
+ data_format=self.data_format,
+ name=f"resample_{i}_{off}",
+ )
+ for i, off in enumerate(self.inputs_offsets)
+ ]
+ self.conv = sep_conv(
+ fpn_num_filters,
+ 3,
+ use_bias=not conv_bn_act_pattern,
+ data_format=self.data_format,
+ name="conv",
+ )
+ self.bn = layers.BatchNormalization(
+ axis=channel_axis(self.data_format), name="bn"
+ )
+
+ def build(self, input_shape):
+ self.edge_weights = None
+ if self.weight_method in ("fastattn", "attn"):
+ self.edge_weights = [
+ self.add_weight(
+ name="WSM" + ("" if i == 0 else f"_{i}"),
+ shape=(),
+ initializer="ones",
+ trainable=True,
+ )
+ for i in range(len(self.inputs_offsets))
+ ]
+ self.built = True
+
+ def fuse(self, nodes):
+ if self.weight_method == "fastattn":
+ weights = [ops.relu(w) for w in self.edge_weights]
+ total = ops.sum(ops.stack(weights)) + 1e-4
+ return ops.sum(
+ ops.stack([n * (w / total) for n, w in zip(nodes, weights)]), axis=0
+ )
+ if self.weight_method == "attn":
+ weights = ops.softmax(ops.stack(self.edge_weights))
+ return ops.sum(
+ ops.stack([n * weights[i] for i, n in enumerate(nodes)]), axis=0
+ )
+ return ops.sum(ops.stack(nodes), axis=0)
+
+ def call(self, feats, target_height, target_width, training=False):
+ nodes = [
+ self.resamples[i](
+ feats[off],
+ target_height=target_height,
+ target_width=target_width,
+ training=training,
+ )
+ for i, off in enumerate(self.inputs_offsets)
+ ]
+ new_node = self.fuse(nodes)
+ if not self.conv_bn_act_pattern:
+ new_node = act(new_node, self.act_type)
+ new_node = self.conv(new_node)
+ new_node = self.bn(new_node, training=training)
+ if self.conv_bn_act_pattern:
+ new_node = act(new_node, self.act_type)
+ return feats + [new_node]
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "feat_level": self.feat_level,
+ "inputs_offsets": self.inputs_offsets,
+ "fpn_num_filters": self.fpn_num_filters,
+ "weight_method": self.weight_method,
+ "act_type": self.act_type,
+ "apply_bn_for_resampling": self.apply_bn_for_resampling,
+ "conv_after_downsample": self.conv_after_downsample,
+ "conv_bn_act_pattern": self.conv_bn_act_pattern,
+ "data_format": self.data_format,
+ }
+ )
+ return config
+
+
+def bifpn_nodes(min_level, max_level):
+ """Google's dynamic BiFPN node topology for ``min_level``..``max_level``.
+
+ Each node lists the offsets (into the growing feature list) it fuses; a
+ top-down path (P{max-1}'..P{min}') then a bottom-up path (P{min+1}''..P{max}'').
+ """
+ num_levels = max_level - min_level + 1
+ node_ids = {min_level + i: [i] for i in range(num_levels)}
+ counter = num_levels
+ nodes = []
+ for level in range(max_level - 1, min_level - 1, -1): # top-down
+ nodes.append(
+ {
+ "feat_level": level,
+ "inputs_offsets": [node_ids[level][-1], node_ids[level + 1][-1]],
+ }
+ )
+ node_ids[level].append(counter)
+ counter += 1
+ for level in range(min_level + 1, max_level + 1): # bottom-up
+ nodes.append(
+ {
+ "feat_level": level,
+ "inputs_offsets": node_ids[level] + [node_ids[level - 1][-1]],
+ }
+ )
+ node_ids[level].append(counter)
+ counter += 1
+ return nodes
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class FPNCells(layers.Layer):
+ """The full BiFPN: ``fpn_cell_repeats`` stacked cells of ``FNode``s over levels
+ ``min_level``..``max_level``. Takes and returns a list of ``max_level-min_level+1``
+ feature maps (all at ``fpn_num_filters`` channels), each level at its own
+ resolution."""
+
+ def __init__(
+ self,
+ min_level,
+ max_level,
+ fpn_num_filters,
+ fpn_cell_repeats,
+ weight_method,
+ act_type,
+ apply_bn_for_resampling,
+ conv_after_downsample,
+ conv_bn_act_pattern,
+ data_format=None,
+ name="fpn_cells",
+ ):
+ super().__init__(name=name)
+ self.min_level = min_level
+ self.max_level = max_level
+ self.fpn_num_filters = fpn_num_filters
+ self.fpn_cell_repeats = fpn_cell_repeats
+ self.weight_method = weight_method
+ self.act_type = act_type
+ self.data_format = get_data_format(data_format)
+ self.nodes_cfg = bifpn_nodes(min_level, max_level)
+ self.cells = []
+ for rep in range(fpn_cell_repeats):
+ fnodes = [
+ FNode(
+ cfg["feat_level"] - min_level,
+ cfg["inputs_offsets"],
+ fpn_num_filters,
+ weight_method,
+ act_type,
+ apply_bn_for_resampling,
+ conv_after_downsample,
+ conv_bn_act_pattern,
+ data_format=self.data_format,
+ name=f"cell_{rep}_fnode{i}",
+ )
+ for i, cfg in enumerate(self.nodes_cfg)
+ ]
+ self.cells.append(fnodes)
+
+ def call(self, feats, level_sizes, training=False):
+ """``level_sizes``: list of (height, width) per level (min_level..max_level)."""
+ for fnodes in self.cells:
+ cur = list(feats)
+ for fnode, cfg in zip(fnodes, self.nodes_cfg):
+ lvl = cfg["feat_level"] - self.min_level
+ th, tw = level_sizes[lvl]
+ cur = fnode(cur, target_height=th, target_width=tw, training=training)
+ # keep the last node produced at each level (reverse-scan, like Google's).
+ feats = []
+ for level in range(self.min_level, self.max_level + 1):
+ for i, cfg in enumerate(reversed(self.nodes_cfg)):
+ if cfg["feat_level"] == level:
+ feats.append(cur[-1 - i])
+ break
+ return feats
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class PredictionHead(layers.Layer):
+ """Shared class or box head: ``repeats`` separable 3x3 convs, each with a
+ per-level BatchNorm and activation (conv weights shared across levels, BN not),
+ then a final separable 3x3 predictor with ``out_channels`` outputs per anchor."""
+
+ def __init__(
+ self,
+ out_channels,
+ num_filters,
+ repeats,
+ num_levels,
+ act_type,
+ head_name,
+ min_level=3,
+ predict_bias_init=0.0,
+ data_format=None,
+ name=None,
+ ):
+ super().__init__(name=name)
+ self.out_channels = out_channels
+ self.num_filters = num_filters
+ self.repeats = repeats
+ self.num_levels = num_levels
+ self.act_type = act_type
+ self.head_name = head_name
+ self.min_level = min_level
+ self.predict_bias_init = predict_bias_init
+ self.data_format = get_data_format(data_format)
+ self.convs = [
+ sep_conv(
+ num_filters,
+ 3,
+ use_bias=True,
+ data_format=self.data_format,
+ name=f"{head_name}-{i}",
+ )
+ for i in range(repeats)
+ ]
+ # one BatchNorm per (repeat, level); conv weights are shared across levels.
+ # BN names carry the absolute pyramid level (min_level + lvl), like Google's.
+ self.bns = [
+ [
+ layers.BatchNormalization(
+ axis=channel_axis(self.data_format),
+ name=f"{head_name}-{i}-bn-{min_level + lvl}",
+ )
+ for lvl in range(num_levels)
+ ]
+ for i in range(repeats)
+ ]
+ self.predict = layers.SeparableConv2D(
+ out_channels,
+ 3,
+ padding="same",
+ depth_multiplier=1,
+ bias_initializer=keras.initializers.Constant(predict_bias_init),
+ data_format=self.data_format,
+ name=f"{head_name}-predict",
+ )
+
+ def call(self, feats, training=False):
+ outputs = []
+ for level_id in range(self.num_levels):
+ x = feats[level_id]
+ for i in range(self.repeats):
+ x = self.convs[i](x)
+ x = self.bns[i][level_id](x, training=training)
+ x = act(x, self.act_type)
+ outputs.append(self.predict(x))
+ return outputs
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "out_channels": self.out_channels,
+ "num_filters": self.num_filters,
+ "repeats": self.repeats,
+ "num_levels": self.num_levels,
+ "act_type": self.act_type,
+ "head_name": self.head_name,
+ "min_level": self.min_level,
+ "predict_bias_init": self.predict_bias_init,
+ "data_format": self.data_format,
+ }
+ )
+ return config
+
+
+def class_predict_bias():
+ """Focal-loss prior bias for the class predictor: ``-log((1 - 0.01) / 0.01)``."""
+ return -math.log((1 - 0.01) / 0.01)
+
+
+def feat_sizes(image_size, max_level):
+ """Per-level (height, width) of the feature pyramid, ceil-halving from the
+ input size (index 0 = input). Matches Google's ``utils.get_feat_sizes``."""
+ if isinstance(image_size, int):
+ h = w = image_size
+ else:
+ h, w = image_size
+ sizes = [(h, w)]
+ for _ in range(1, max_level + 1):
+ h, w = (h - 1) // 2 + 1, (w - 1) // 2 + 1
+ sizes.append((h, w))
+ return sizes
+
+
+def generate_anchor_boxes(
+ min_level, max_level, num_scales, aspect_ratios, anchor_scale, image_size
+):
+ """Multiscale anchor boxes as ``(N, 4)`` in ``[ymin, xmin, ymax, xmax]`` pixel
+ coordinates, ordered position-major then (octave, aspect)-minor per level, then
+ concatenated over levels. A direct port of Google AutoML ``anchors.Anchors``."""
+ sizes = feat_sizes(image_size, max_level)
+ if isinstance(anchor_scale, (list, tuple)):
+ anchor_scales = list(anchor_scale)
+ else:
+ anchor_scales = [anchor_scale] * (max_level - min_level + 1)
+
+ boxes_all = []
+ for level in range(min_level, max_level + 1):
+ feat_h, feat_w = sizes[level]
+ stride_h = sizes[0][0] / float(feat_h)
+ stride_w = sizes[0][1] / float(feat_w)
+ cx = (ops.arange(feat_w, dtype="float32") + 0.5) * stride_w
+ cy = (ops.arange(feat_h, dtype="float32") + 0.5) * stride_h
+ xv, yv = ops.meshgrid(cx, cy, indexing="xy")
+ xv, yv = ops.reshape(xv, (-1,)), ops.reshape(yv, (-1,))
+ boxes_level = []
+ for scale_octave in range(num_scales):
+ for aspect in aspect_ratios:
+ octave = 2 ** (scale_octave / float(num_scales))
+ base_x = anchor_scales[level - min_level] * stride_w * octave
+ base_y = anchor_scales[level - min_level] * stride_h * octave
+ aspect_x = math.sqrt(aspect)
+ aspect_y = 1.0 / aspect_x
+ half_x = base_x * aspect_x / 2.0
+ half_y = base_y * aspect_y / 2.0
+ boxes = ops.stack(
+ [yv - half_y, xv - half_x, yv + half_y, xv + half_x], axis=-1
+ )
+ boxes_level.append(ops.expand_dims(boxes, axis=1))
+ boxes_level = ops.concatenate(boxes_level, axis=1)
+ boxes_all.append(ops.reshape(boxes_level, (-1, 4)))
+ return ops.cast(ops.concatenate(boxes_all, axis=0), "float32")
+
+
+def decode_box_outputs(pred_boxes, anchor_boxes):
+ """Invert the anchor box regression: ``(ty, tx, th, tw)`` relative to anchors ->
+ absolute ``[ymin, xmin, ymax, xmax]``. ``anchor_boxes`` are ``[ymin, xmin, ymax,
+ xmax]``. Matches Google AutoML ``anchors.decode_box_outputs``."""
+ ycenter_a = (anchor_boxes[..., 0] + anchor_boxes[..., 2]) / 2
+ xcenter_a = (anchor_boxes[..., 1] + anchor_boxes[..., 3]) / 2
+ ha = anchor_boxes[..., 2] - anchor_boxes[..., 0]
+ wa = anchor_boxes[..., 3] - anchor_boxes[..., 1]
+ ty, tx, th, tw = (
+ pred_boxes[..., 0],
+ pred_boxes[..., 1],
+ pred_boxes[..., 2],
+ pred_boxes[..., 3],
+ )
+ w = ops.exp(tw) * wa
+ h = ops.exp(th) * ha
+ ycenter = ty * ha + ycenter_a
+ xcenter = tx * wa + xcenter_a
+ return ops.stack(
+ [
+ ycenter - h / 2.0,
+ xcenter - w / 2.0,
+ ycenter + h / 2.0,
+ xcenter + w / 2.0,
+ ],
+ axis=-1,
+ )
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class DecodeBoxes(layers.Layer):
+ """Turn raw per-anchor box regressions ``(B, N, 4)`` into absolute
+ ``[ymin, xmin, ymax, xmax]`` boxes using pre-generated anchors held as a
+ non-trainable constant."""
+
+ def __init__(
+ self,
+ min_level,
+ max_level,
+ num_scales,
+ aspect_ratios,
+ anchor_scale,
+ image_size,
+ name="decode_boxes",
+ **kwargs,
+ ):
+ super().__init__(name=name, **kwargs)
+ self.min_level = min_level
+ self.max_level = max_level
+ self.num_scales = num_scales
+ self.aspect_ratios = tuple(aspect_ratios)
+ self.anchor_scale = anchor_scale
+ self.image_size = image_size
+ self.anchors = generate_anchor_boxes(
+ min_level, max_level, num_scales, aspect_ratios, anchor_scale, image_size
+ )
+
+ def build(self, input_shape):
+ self.built = True
+
+ def call(self, box_outputs):
+ # Anchors are a fixed function of the config, so they are baked into the graph
+ # as a constant rather than stored as a weight. This keeps EfficientDetDetect's
+ # weight set identical to EfficientDetModel's, so both load one hosted file.
+ return decode_box_outputs(
+ box_outputs, ops.cast(self.anchors, box_outputs.dtype)
+ )
+
+ def compute_output_shape(self, input_shape):
+ return input_shape
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "min_level": self.min_level,
+ "max_level": self.max_level,
+ "num_scales": self.num_scales,
+ "aspect_ratios": self.aspect_ratios,
+ "anchor_scale": self.anchor_scale,
+ "image_size": self.image_size,
+ }
+ )
+ return config
+
+
+def iou_against(box, boxes):
+ """IoU of one ``[ymin, xmin, ymax, xmax]`` box against an array of boxes."""
+ ymin = ops.maximum(box[0], boxes[:, 0])
+ xmin = ops.maximum(box[1], boxes[:, 1])
+ ymax = ops.minimum(box[2], boxes[:, 2])
+ xmax = ops.minimum(box[3], boxes[:, 3])
+ inter = ops.maximum(0.0, ymax - ymin) * ops.maximum(0.0, xmax - xmin)
+ area = (box[2] - box[0]) * (box[3] - box[1])
+ areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+ return inter / (area + areas - inter + 1e-8)
+
+
+def greedy_nms(boxes, scores, iou_threshold, score_threshold, max_output):
+ """Greedy hard NMS. Walks boxes in descending score order, keeps one and suppresses
+ everything it overlaps by more than ``iou_threshold``, and stops once the score drops
+ to ``score_threshold`` (scores are sorted, so nothing later can clear it). Returns the
+ kept indices, highest score first."""
+ order = ops.argsort(-scores)
+ total = int(order.shape[0])
+ suppressed = ops.zeros((total,), dtype="bool")
+ keep = []
+ for rank in range(total):
+ idx = int(order[rank])
+ if float(scores[idx]) <= score_threshold:
+ break
+ if bool(suppressed[idx]):
+ continue
+ keep.append(idx)
+ if len(keep) >= max_output:
+ break
+ suppressed = ops.logical_or(
+ suppressed, iou_against(boxes[idx], boxes) > iou_threshold
+ )
+ return keep
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class EfficientDetNMS(layers.Layer):
+ """Hard non-max suppression over decoded detections.
+
+ Given decoded boxes ``(B, N, 4)`` in ``[ymin, xmin, ymax, xmax]`` and per-class
+ sigmoid scores ``(B, N, num_classes)``, returns a padded ``(B, max_detections, 6)``
+ tensor of ``[ymin, xmin, ymax, xmax, score, class_id]`` (rows past the detection
+ count are zero).
+
+ With ``class_agnostic=True`` (the default, matching Google AutoML's
+ ``postprocess_global``) each anchor keeps only its single highest-scoring class
+ and one NMS runs across all classes together, so a single object cannot yield two
+ boxes under different labels (e.g. a dog also reported as a cat). With
+ ``class_agnostic=False`` NMS runs independently per class
+ (``postprocess_per_class``). Runs eagerly (detection post-processing), applied
+ outside the symbolic graph."""
+
+ def __init__(
+ self,
+ iou_threshold=0.5,
+ score_threshold=0.05,
+ max_detections=100,
+ pre_nms_top_k=5000,
+ class_agnostic=True,
+ name="nms",
+ **kwargs,
+ ):
+ super().__init__(name=name, **kwargs)
+ self.iou_threshold = iou_threshold
+ self.score_threshold = score_threshold
+ self.max_detections = max_detections
+ self.pre_nms_top_k = pre_nms_top_k
+ self.class_agnostic = class_agnostic
+
+ def top_k_prefilter(self, boxes, scores):
+ """Keep the ``pre_nms_top_k`` anchors with the highest best-class score."""
+ best = ops.max(scores, axis=-1)
+ if int(best.shape[0]) > self.pre_nms_top_k:
+ top = ops.top_k(best, self.pre_nms_top_k)[1]
+ return ops.take(boxes, top, axis=0), ops.take(scores, top, axis=0)
+ return boxes, scores
+
+ def detections_global(self, boxes, scores):
+ """Class-agnostic: one best class per anchor, then a single NMS."""
+ class_ids = ops.argmax(scores, axis=-1)
+ class_scores = ops.max(scores, axis=-1)
+ keep = greedy_nms(
+ boxes,
+ class_scores,
+ self.iou_threshold,
+ self.score_threshold,
+ self.max_detections,
+ )
+ return [
+ (
+ boxes[k][0],
+ boxes[k][1],
+ boxes[k][2],
+ boxes[k][3],
+ class_scores[k],
+ ops.cast(class_ids[k], "float32"),
+ )
+ for k in keep
+ ]
+
+ def detections_per_class(self, boxes, scores):
+ """Independent NMS per class."""
+ dets = []
+ for c in range(int(scores.shape[-1])):
+ cls_scores = scores[:, c]
+ keep = greedy_nms(
+ boxes,
+ cls_scores,
+ self.iou_threshold,
+ self.score_threshold,
+ self.max_detections,
+ )
+ dets.extend(
+ (
+ boxes[k][0],
+ boxes[k][1],
+ boxes[k][2],
+ boxes[k][3],
+ cls_scores[k],
+ float(c),
+ )
+ for k in keep
+ )
+ return dets
+
+ def pad_detections(self, dets):
+ """Sort by score and pad the detection tuples to ``(max_detections, 6)``."""
+ if not dets:
+ return ops.zeros((self.max_detections, 6), dtype="float32")
+ dets = sorted(dets, key=lambda d: float(d[4]), reverse=True)
+ dets = dets[: self.max_detections]
+ rows = ops.stack(
+ [ops.stack([ops.cast(v, "float32") for v in d]) for d in dets], axis=0
+ )
+ pad = self.max_detections - int(rows.shape[0])
+ if pad > 0:
+ rows = ops.concatenate([rows, ops.zeros((pad, 6), dtype="float32")], axis=0)
+ return rows
+
+ def call(self, boxes, scores):
+ boxes = ops.convert_to_tensor(boxes)
+ scores = ops.convert_to_tensor(scores)
+ results = []
+ for b in range(int(boxes.shape[0])):
+ bx, sc = self.top_k_prefilter(boxes[b], scores[b])
+ if self.class_agnostic:
+ dets = self.detections_global(bx, sc)
+ else:
+ dets = self.detections_per_class(bx, sc)
+ results.append(self.pad_detections(dets))
+ return ops.stack(results, axis=0)
+
+ def compute_output_shape(self, boxes_shape, scores_shape):
+ return (boxes_shape[0], self.max_detections, 6)
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "iou_threshold": self.iou_threshold,
+ "score_threshold": self.score_threshold,
+ "max_detections": self.max_detections,
+ "pre_nms_top_k": self.pre_nms_top_k,
+ "class_agnostic": self.class_agnostic,
+ }
+ )
+ return config
diff --git a/zeromodels/models/efficientdet/efficientdet_model.py b/zeromodels/models/efficientdet/efficientdet_model.py
new file mode 100644
index 00000000..90cec73d
--- /dev/null
+++ b/zeromodels/models/efficientdet/efficientdet_model.py
@@ -0,0 +1,331 @@
+import keras
+from keras import layers
+
+from zeromodels.base import BaseModel, CheckpointSource
+from zeromodels.models.efficientnet.efficientnet_model import (
+ efficientnet_backbone_feature,
+)
+from zeromodels.utils.image_util import get_data_format
+
+from .efficientdet_config import EfficientDetConfig
+from .efficientdet_layers import (
+ DecodeBoxes,
+ EfficientDetResample,
+ FPNCells,
+ PredictionHead,
+ channel_axis,
+ class_predict_bias,
+)
+
+# EfficientNet (width, depth, dropout) coefficients per EfficientDet backbone.
+BACKBONE_COEFFS = {
+ "efficientnet_b0": (1.0, 1.0, 0.2),
+ "efficientnet_b1": (1.0, 1.1, 0.2),
+ "efficientnet_b2": (1.1, 1.2, 0.3),
+ "efficientnet_b3": (1.2, 1.4, 0.3),
+ "efficientnet_b4": (1.4, 1.8, 0.4),
+ "efficientnet_b5": (1.6, 2.2, 0.4),
+ "efficientnet_b6": (1.8, 2.6, 0.5),
+}
+
+CONFIG_ATTRS = (
+ "backbone_name",
+ "image_size",
+ "num_classes",
+ "min_level",
+ "max_level",
+ "num_scales",
+ "aspect_ratios",
+ "anchor_scale",
+ "fpn_num_filters",
+ "fpn_cell_repeats",
+ "box_class_repeats",
+ "act_type",
+ "separable_conv",
+ "apply_bn_for_resampling",
+ "conv_after_downsample",
+ "conv_bn_act_pattern",
+ "fpn_weight_method",
+ "survival_prob",
+)
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class EfficientDetModel(BaseModel):
+ """EfficientDet backbone + BiFPN + heads (raw outputs, Google AutoML).
+
+ An EfficientNet backbone (multi-level features P3-P5) is extended with strided
+ resamples to P6-P7, fused by a repeated weighted BiFPN, and read out by shared
+ class + box heads. Takes an ``images`` input ``(B, image_size, image_size, 3)``
+ and returns ``{"class_outputs", "box_outputs"}``: lists of per-level tensors
+ ``(B, H_l, W_l, num_anchors * num_classes)`` and ``(B, H_l, W_l, num_anchors * 4)``
+ for levels ``min_level``..``max_level``. Decode with anchors + NMS.
+
+ References:
+ - [EfficientDet: Scalable and Efficient Object Detection](https://arxiv.org/abs/1911.09070)
+
+ Args:
+ See :class:`EfficientDetConfig`. Defaults describe EfficientDet-D0.
+ name: String, model name. Defaults to `"EfficientDet"`.
+
+ Returns:
+ A Keras `Model` instance.
+ """
+
+ BASE_WEIGHT_CONFIG = None
+ HF_MODEL_TYPE = "efficientdet"
+ config_class = EfficientDetConfig
+
+ def __init__(
+ self,
+ backbone_name="efficientnet_b0",
+ image_size=512,
+ num_classes=90,
+ min_level=3,
+ max_level=7,
+ num_scales=3,
+ aspect_ratios=(1.0, 2.0, 0.5),
+ anchor_scale=4.0,
+ fpn_num_filters=64,
+ fpn_cell_repeats=3,
+ box_class_repeats=3,
+ act_type="swish",
+ separable_conv=True,
+ apply_bn_for_resampling=True,
+ conv_after_downsample=False,
+ conv_bn_act_pattern=False,
+ fpn_weight_method="fastattn",
+ survival_prob=None,
+ name="EfficientDetModel",
+ **kwargs,
+ ):
+ for k in ("model", "hf_id", "url", "num_classes_"):
+ kwargs.pop(k, None)
+
+ data_format = get_data_format()
+ num_anchors = num_scales * len(aspect_ratios)
+ width, depth, dropout = BACKBONE_COEFFS[backbone_name]
+ level_sizes = [
+ (image_size // (2**level), image_size // (2**level))
+ for level in range(min_level, max_level + 1)
+ ]
+
+ input_shape = (
+ (3, image_size, image_size)
+ if data_format == "channels_first"
+ else (image_size, image_size, 3)
+ )
+ image = layers.Input(shape=input_shape, name="images")
+
+ # Backbone P3/P4/P5 (strides 8/16/32 -> stages[2:5]).
+ stages = efficientnet_backbone_feature(
+ image,
+ width_coefficient=width,
+ depth_coefficient=depth,
+ dropout_rate=dropout,
+ data_format=data_format,
+ channels_axis=channel_axis(data_format),
+ return_stages=True,
+ )
+ feats = [stages[2], stages[3], stages[4]]
+
+ # Coarser levels P6..max_level by strided resampling of the last feature.
+ for level in range(6, max_level + 1):
+ th, tw = level_sizes[level - min_level]
+ feats.append(
+ EfficientDetResample(
+ fpn_num_filters,
+ apply_bn=apply_bn_for_resampling,
+ conv_after_downsample=conv_after_downsample,
+ data_format=data_format,
+ name=f"resample_p{level}",
+ )(feats[-1], target_height=th, target_width=tw)
+ )
+
+ fpn_feats = FPNCells(
+ min_level,
+ max_level,
+ fpn_num_filters,
+ fpn_cell_repeats,
+ fpn_weight_method,
+ act_type,
+ apply_bn_for_resampling,
+ conv_after_downsample,
+ conv_bn_act_pattern,
+ data_format=data_format,
+ name="fpn_cells",
+ )(feats, level_sizes=level_sizes)
+
+ num_levels = max_level - min_level + 1
+ class_outputs = PredictionHead(
+ num_anchors * num_classes,
+ fpn_num_filters,
+ box_class_repeats,
+ num_levels,
+ act_type,
+ "class",
+ min_level=min_level,
+ predict_bias_init=class_predict_bias(),
+ data_format=data_format,
+ name="class_net",
+ )(fpn_feats)
+ box_outputs = PredictionHead(
+ num_anchors * 4,
+ fpn_num_filters,
+ box_class_repeats,
+ num_levels,
+ act_type,
+ "box",
+ min_level=min_level,
+ predict_bias_init=0.0,
+ data_format=data_format,
+ name="box_net",
+ )(fpn_feats)
+
+ super().__init__(
+ inputs=image,
+ outputs={"class_outputs": class_outputs, "box_outputs": box_outputs},
+ name=name,
+ **kwargs,
+ )
+ for attr in CONFIG_ATTRS:
+ setattr(self, attr, locals()[attr])
+
+ def get_config(self):
+ config = super().get_config()
+ config.update({k: getattr(self, k) for k in CONFIG_ATTRS})
+ config["name"] = self.name
+ return config
+
+ @classmethod
+ def from_config(cls, config):
+ return cls(**config)
+
+
+@keras.saving.register_keras_serializable(package="zeromodels")
+class EfficientDetDetect(BaseModel):
+ """EfficientDet object detector: :class:`EfficientDetModel` + anchor decoding.
+
+ Wraps the backbone/BiFPN/heads model and turns its raw per-level outputs into
+ decoded detections. The functional graph flattens and concatenates the per-level
+ class and box outputs, applies sigmoid to the class logits, and decodes the box
+ regressions against pre-generated anchors, exposing
+ ``{"boxes": (B, N, 4), "scores": (B, N, num_classes)}`` (boxes are
+ ``[ymin, xmin, ymax, xmax]`` in ``image_size`` pixel coordinates). Non-max
+ suppression and rescaling to the original image are applied afterwards by
+ :meth:`EfficientDetImageProcessor.post_process_object_detection`.
+
+ References:
+ - [EfficientDet: Scalable and Efficient Object Detection](https://arxiv.org/abs/1911.09070)
+
+ Args:
+ See :class:`EfficientDetConfig`. Defaults describe EfficientDet-D0.
+ name: String, model name. Defaults to `"EfficientDetDetect"`.
+
+ Returns:
+ A Keras `Model` instance.
+ """
+
+ BASE_WEIGHT_CONFIG = None
+ HF_MODEL_TYPE = "efficientdet"
+ config_class = EfficientDetConfig
+ # EfficientDetDetect shares its weights with EfficientDetModel (identical backbone,
+ # BiFPN and heads; decoding adds no weights). Hosted repos declare the canonical
+ # EfficientDetModel, and this class loads that same file by copying weights out.
+ HUB_REPO_SIBLINGS = frozenset({"EfficientDetModel"})
+ CHECKPOINT_SOURCE = CheckpointSource("EfficientDetModel")
+
+ def __init__(
+ self,
+ backbone_name="efficientnet_b0",
+ image_size=512,
+ num_classes=90,
+ min_level=3,
+ max_level=7,
+ num_scales=3,
+ aspect_ratios=(1.0, 2.0, 0.5),
+ anchor_scale=4.0,
+ fpn_num_filters=64,
+ fpn_cell_repeats=3,
+ box_class_repeats=3,
+ act_type="swish",
+ separable_conv=True,
+ apply_bn_for_resampling=True,
+ conv_after_downsample=False,
+ conv_bn_act_pattern=False,
+ fpn_weight_method="fastattn",
+ survival_prob=None,
+ name="EfficientDetDetect",
+ **kwargs,
+ ):
+ for k in ("model", "hf_id", "url", "num_classes_"):
+ kwargs.pop(k, None)
+
+ base = EfficientDetModel(
+ backbone_name=backbone_name,
+ image_size=image_size,
+ num_classes=num_classes,
+ min_level=min_level,
+ max_level=max_level,
+ num_scales=num_scales,
+ aspect_ratios=aspect_ratios,
+ anchor_scale=anchor_scale,
+ fpn_num_filters=fpn_num_filters,
+ fpn_cell_repeats=fpn_cell_repeats,
+ box_class_repeats=box_class_repeats,
+ act_type=act_type,
+ separable_conv=separable_conv,
+ apply_bn_for_resampling=apply_bn_for_resampling,
+ conv_after_downsample=conv_after_downsample,
+ conv_bn_act_pattern=conv_bn_act_pattern,
+ fpn_weight_method=fpn_weight_method,
+ survival_prob=survival_prob,
+ name=f"{name}_model",
+ )
+ class_outputs = base.output["class_outputs"]
+ box_outputs = base.output["box_outputs"]
+
+ # Flatten each level's head output to (B, H*W*anchors, last). Under
+ # channels_first the head output is (B, anchors*last, H, W); transpose it to
+ # channels_last first so the anchor ordering (position-major) matches the
+ # anchor grid regardless of data format -> identical outputs either way.
+ data_format = get_data_format()
+
+ def flatten_levels(tensors, last, name):
+ flats = []
+ for tensor in tensors:
+ if data_format == "channels_first":
+ tensor = layers.Permute((2, 3, 1))(tensor)
+ flats.append(layers.Reshape((-1, last))(tensor))
+ return layers.Concatenate(axis=1, name=name)(flats)
+
+ cls_flat = flatten_levels(class_outputs, num_classes, "class_concat")
+ box_flat = flatten_levels(box_outputs, 4, "box_concat")
+ scores = layers.Activation("sigmoid", name="scores")(cls_flat)
+ boxes = DecodeBoxes(
+ min_level,
+ max_level,
+ num_scales,
+ aspect_ratios,
+ anchor_scale,
+ image_size,
+ )(box_flat)
+
+ super().__init__(
+ inputs=base.input,
+ outputs={"boxes": boxes, "scores": scores},
+ name=name,
+ **kwargs,
+ )
+ for attr in CONFIG_ATTRS:
+ setattr(self, attr, locals()[attr])
+
+ def get_config(self):
+ config = super().get_config()
+ config.update({k: getattr(self, k) for k in CONFIG_ATTRS})
+ config["name"] = self.name
+ return config
+
+ @classmethod
+ def from_config(cls, config):
+ return cls(**config)
diff --git a/zeromodels/models/granite_speech/convert_granite_speech_hf_to_keras.py b/zeromodels/models/granite_speech/convert_granite_speech_hf_to_keras.py
index d180dae0..bb2646f9 100644
--- a/zeromodels/models/granite_speech/convert_granite_speech_hf_to_keras.py
+++ b/zeromodels/models/granite_speech/convert_granite_speech_hf_to_keras.py
@@ -58,9 +58,6 @@ def transfer_granite_speech_weights(keras_model, hf_state_dict):
for weight in tqdm(keras_model.weights, desc="Transferring weights to Keras"):
keras_dotted = weight.path.replace("/", ".")
- # The conformer encoder + Q-Former projector run inside the GraniteSpeechAudioFeatures
- # wrapper, so keras paths them under "audio_features/"; the HF checkpoint keeps them at
- # the top level (encoder.* / projector.*). Strip the wrapper prefix before mapping.
if keras_dotted.startswith("audio_features."):
keras_dotted = keras_dotted[len("audio_features.") :]
if keras_dotted == "token_embedding.embeddings":