Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,220 @@
+//! Region Proposal Network (RPN) and anchor generation.
+
+use crate::utils::nms;
+use crate::{BoundingBox, VisionError, VisionResult};
+use rtx_tensor::{Device, Tensor};
+use tracing::debug;
+
+use super::layers::ConvBlock;
+
+/// Region Proposal Network (RPN)
+pub struct RegionProposalNetwork {
+    conv: ConvBlock,
+    cls_logits: ConvBlock,
+    bbox_pred: ConvBlock,
+    anchor_generator: AnchorGenerator,
+    anchor_matcher: AnchorMatcher,
+}
+
+impl RegionProposalNetwork {
+    pub fn new(in_channels: usize, num_anchors_per_location: usize) -> VisionResult<Self> {
+        let conv = ConvBlock::new(in_channels, in_channels, 3, 1, 1)?;
+        let cls_logits = ConvBlock::new(in_channels, num_anchors_per_location, 1, 1, 0)?;
+        let bbox_pred = ConvBlock::new(in_channels, num_anchors_per_location * 4, 1, 1, 0)?;
+
+        let anchor_generator = AnchorGenerator::new(vec![8.0, 16.0, 32.0], vec![0.5, 1.0, 2.0])?;
+        let anchor_matcher = AnchorMatcher::new(0.3, 0.7)?;
+
+        Ok(Self {
+            conv,
+            cls_logits,
+            bbox_pred,
+            anchor_generator,
+            anchor_matcher,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        features: &[Tensor],
+        image_sizes: &[(usize, usize)],
+    ) -> VisionResult<(Vec<Tensor>, Vec<Tensor>, Vec<Tensor>)> {
+        let mut objectness_logits = Vec::new();
+        let mut pred_bbox_deltas = Vec::new();
+        let mut anchors = Vec::new();
+
+        for (feature, &image_size) in features.iter().zip(image_sizes.iter()) {
+            // RPN head
+            let conv_out = self.conv.forward(feature)?;
+
+            // Classification (objectness)
+            let cls_logits = self.cls_logits.forward(&conv_out)?;
+            objectness_logits.push(cls_logits);
+
+            // Regression (bbox deltas)
+            let bbox_deltas = self.bbox_pred.forward(&conv_out)?;
+            pred_bbox_deltas.push(bbox_deltas);
+
+            // Generate anchors
+            let feature_shape = (feature.shape()[2], feature.shape()[3]);
+            let level_anchors = self
+                .anchor_generator
+                .generate_anchors(feature_shape, image_size)?;
+            anchors.push(level_anchors);
+        }
+
+        Ok((objectness_logits, pred_bbox_deltas, anchors))
+    }
+
+    pub fn filter_proposals(
+        &self,
+        proposals: Vec<BoundingBox>,
+        _objectness_logits: &Tensor,
+        _image_size: (usize, usize),
+        training: bool,
+    ) -> VisionResult<Vec<BoundingBox>> {
+        let num_proposals_train = 2000;
+        let num_proposals_test = 1000;
+        let nms_threshold = 0.7;
+
+        let max_proposals = if training {
+            num_proposals_train
+        } else {
+            num_proposals_test
+        };
+
+        // Apply NMS to proposals
+        let mut filtered_proposals = proposals;
+        filtered_proposals.retain(|prop| prop.confidence > 0.5);
+
+        if filtered_proposals.len() > max_proposals {
+            filtered_proposals.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
+            filtered_proposals.truncate(max_proposals);
+        }
+
+        // Apply NMS
+        let final_proposals = nms::apply_nms(filtered_proposals, nms_threshold)?;
+
+        debug!("RPN generated {} proposals", final_proposals.len());
+        Ok(final_proposals)
+    }
+}
+
+/// Anchor generator for RPN
+pub struct AnchorGenerator {
+    scales: Vec<f32>,
+    aspect_ratios: Vec<f32>,
+    pub base_anchors: Vec<BoundingBox>,
+}
+
+impl AnchorGenerator {
+    pub fn new(scales: Vec<f32>, aspect_ratios: Vec<f32>) -> VisionResult<Self> {
+        let mut base_anchors = Vec::new();
+
+        for &scale in &scales {
+            for &ratio in &aspect_ratios {
+                let width = scale * ratio.sqrt();
+                let height = scale / ratio.sqrt();
+
+                base_anchors.push(BoundingBox::new(
+                    -width / 2.0,
+                    -height / 2.0,
+                    width,
+                    height,
+                    1.0, // Default confidence
+                    0,   // Default class
+                ));
+            }
+        }
+
+        Ok(Self {
+            scales,
+            aspect_ratios,
+            base_anchors,
+        })
+    }
+
+    pub fn generate_anchors(
+        &self,
+        feature_shape: (usize, usize),
+        image_size: (usize, usize),
+    ) -> VisionResult<Tensor> {
+        let (feat_h, feat_w) = feature_shape;
+        let (img_h, img_w) = image_size;
+
+        let stride_h = img_h as f32 / feat_h as f32;
+        let stride_w = img_w as f32 / feat_w as f32;
+
+        let mut all_anchors = Vec::new();
+
+        for y in 0..feat_h {
+            for x in 0..feat_w {
+                let center_x = (x as f32 + 0.5) * stride_w;
+                let center_y = (y as f32 + 0.5) * stride_h;
+
+                for base_anchor in &self.base_anchors {
+                    let anchor_x = center_x + base_anchor.x;
+                    let anchor_y = center_y + base_anchor.y;
+
+                    all_anchors.push([
+                        anchor_x,
+                        anchor_y,
+                        anchor_x + base_anchor.width,
+                        anchor_y + base_anchor.height,
+                    ]);
+                }
+            }
+        }
+
+        // Convert Vec<[f32; 4]> to Tensor
+        let flattened: Vec<f32> = all_anchors.iter().flat_map(|a| a.iter().copied()).collect();
+        Tensor::from_vec(flattened, &[all_anchors.len(), 4], &Device::default())
+            .map_err(|e| VisionError::tensor_error_with_source("Failed to create anchor tensor", e))
+    }
+}
+
+/// Anchor matcher for training
+pub struct AnchorMatcher {
+    low_threshold: f32,
+    high_threshold: f32,
+}
+
+impl AnchorMatcher {
+    pub fn new(low_threshold: f32, high_threshold: f32) -> VisionResult<Self> {
+        Ok(Self {
+            low_threshold,
+            high_threshold,
+        })
+    }
+
+    pub fn match_anchors(
+        &self,
+        anchors: &[BoundingBox],
+        targets: &[BoundingBox],
+    ) -> VisionResult<Vec<i32>> {
+        let mut matches = vec![-1i32; anchors.len()]; // -1: ignore, 0: negative, >0: positive (target idx + 1)
+
+        for (anchor_idx, anchor) in anchors.iter().enumerate() {
+            let mut best_iou = 0.0;
+            let mut best_target_idx = None;
+
+            for (target_idx, target) in targets.iter().enumerate() {
+                let iou = anchor.iou(target);
+
+                if iou > best_iou {
+                    best_iou = iou;
+                    best_target_idx = Some(target_idx);
+                }
+            }
+
+            if best_iou < self.low_threshold {
+                matches[anchor_idx] = 0; // Negative
+            } else if best_iou >= self.high_threshold {
+                matches[anchor_idx] = (best_target_idx.unwrap() + 1) as i32; // Positive
+            }
+            // else: ignore (keep -1)
+        }
+
+        Ok(matches)
+    }
+}