rustytorch/crates/models/rtx-vision-advanced/src/detection/rcnn/rpn.rs

//! 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)
    }
}