Phát Hiện Đối Tượng và Phân Đoạn

Phát hiện đối tượng (object detection) định vị và phân loại mọi vật thể trong ảnh; phân đoạn (segmentation) gán nhãn cho từng pixel. File này đề cập đến IoU, mAP, anchor box, họ R-CNN, YOLO, SSD, Mạng Kim Tự Tháp Đặc Trưng (FPN), phân đoạn ngữ nghĩa/thực thể/panoptic (U-Net, Mask R-CNN, SAM), và các metric đánh giá chúng.

• Phân loại ảnh (file 02) trả lời "có gì trong ảnh này?" Phát hiện đối tượng đặt một câu hỏi khó hơn: "có những vật thể nào trong ảnh này, và chúng ở đâu?"

• Phân đoạn đi xa hơn nữa: "pixel nào thuộc về vật thể hoặc danh mục nào?" Các tác vụ này tạo thành một hệ phân cấp về độ hiểu biết không gian ngày càng chính xác.

• Một mô hình phát hiện đối tượng (object detection) đầu ra một tập các khung bao (bounding box), mỗi khung được định nghĩa bởi bốn toạ độ (góc trên bên trái \(x, y\), chiều rộng, chiều cao) cùng với nhãn lớp và điểm số tin cậy (confidence score). Một ảnh đơn có thể chứa không, một, hoặc hàng trăm vật thể từ nhiều lớp khác nhau.

• Giao trên hợp (Intersection over Union - IoU) đo lường mức độ khớp giữa khung bao dự đoán và khung bao thực tế (ground truth). Nó là diện tích phần giao chia cho diện tích phần hợp:

\[ \text{IoU} = \frac{\text{Diện tích giao}}{\text{Diện tích hợp}} \]

• IoU bằng 1 nghĩa là chồng khớp hoàn hảo; IoU bằng 0 nghĩa là không có chồng lấn nào. Ngưỡng tiêu chuẩn cho một phát hiện "đúng" là IoU \(\geq 0.5\), mặc dù các ngưỡng chặt hơn (0.75, 0.9) cũng được sử dụng.

• Một phát hiện là dương tính thật (true positive - TP) nếu IoU của nó với một khung bao ground truth vượt quá ngưỡng và lớp đúng.

• Dương tính giả (false positive - FP) là một khung bao dự đoán không khớp với bất kỳ ground truth nào.

• Âm tính giả (false negative - FN) là một vật thể ground truth mà không có dự đoán nào khớp. Đây là các khái niệm precision/recall tương tự từ chương 06.

• Độ chính xác trung bình (Average Precision - AP) tóm tắt chất lượng phát hiện cho một lớp. Với mỗi lớp, xếp hạng tất cả các phát hiện theo điểm số tin cậy, tính precision và recall tại mỗi hạng, và tính diện tích dưới đường cong precision-recall:

\[ \text{AP} = \int_0^1 p(r) \, dr \]

• Trong thực tế, đường cong được nội suy: tại mỗi mức recall, precision được đặt thành precision tối đa tại bất kỳ recall \(\geq r\). Điều này làm mượt đường cong và làm nó giảm đơn điệu.

• Trung bình độ chính xác trung bình (Mean Average Precision - mAP) tính trung bình AP trên tất cả các lớp. "mAP@0.5" sử dụng ngưỡng IoU 0.5. "mAP@[.5:.95]" (tiêu chuẩn COCO) tính trung bình mAP trên mười ngưỡng IoU từ 0.5 đến 0.95 với bước 0.05, thưởng cho cả phát hiện và định vị chính xác.

• Không có cực đại giả (Non-Maximum Suppression - NMS) loại bỏ các phát hiện trùng lặp. Khi một mô hình dự đoán nhiều khung bao chồng lấn cho cùng một vật thể, NMS giữ khung bao có điểm số tin cậy cao nhất và loại bỏ tất cả các khung khác chồng lấn với nó trên một ngưỡng IoU. Điều này được áp dụng cho từng lớp sau khi mô hình tạo ra các dự đoán thô.

• Bộ phát hiện hai giai đoạn (two-stage detectors) đầu tiên đề xuất các vùng ứng viên, sau đó phân loại và tinh chỉnh từng đề xuất.

• R-CNN (Girshick và cộng sự, 2014) là bộ phát hiện học sâu thành công đầu tiên. Nó sử dụng selective search (một thuật toán cổ điển) để đề xuất ~2,000 vùng ứng viên, biến dạng mỗi vùng về kích thước cố định, chạy từng vùng qua một CNN độc lập, và phân loại với SVM (chương 06). R-CNN chính xác nhưng cực kỳ chậm: nó chạy CNN 2,000 lần mỗi ảnh.

• Fast R-CNN (Girshick, 2015) giải quyết sự dư thừa bằng cách chạy CNN một lần trên toàn bộ ảnh để tạo ra một feature map dùng chung, sau đó trích xuất đặc trưng cho mỗi đề xuất từ feature map dùng chung đó bằng RoI pooling (Region of Interest pooling).

• RoI pooling lấy một vùng có kích thước thay đổi của feature map và tạo ra đầu ra kích thước cố định bằng cách chia vùng đó thành một lưới và max-pooling trong mỗi ô. Cách này nhanh hơn nhiều vì tính toán CNN đắt đỏ chỉ xảy ra một lần.

• Faster R-CNN (Ren và cộng sự, 2015) loại bỏ thuật toán đề xuất vùng bên ngoài bằng cách giới thiệu Mạng Đề xuất Vùng (Region Proposal Network - RPN), một CNN nhỏ chạy trên feature map dùng chung và dự đoán các đề xuất trực tiếp. RPN trượt một cửa sổ nhỏ trên feature map và, tại mỗi vị trí, dự đoán \(k\) đề xuất (một cho mỗi anchor box).

• Anchor box là các khung bao được định nghĩa trước tại mỗi vị trí không gian của feature map, bao phủ các tỷ lệ và tỷ lệ khung hình khác nhau (ví dụ: ba tỷ lệ \(\times\) ba tỷ lệ khung hình = 9 anchor mỗi vị trí). RPN dự đoán hai thứ cho mỗi anchor: điểm số đối tượng (object vs background) và độ dịch chuyển toạ độ để tinh chỉnh anchor thành một đề xuất chặt chẽ hơn. Tham số hoá này làm bài toán hồi quy dễ hơn: thay vì dự đoán toạ độ tuyệt đối, mạng dự đoán các điều chỉnh nhỏ so với một khung xuất phát hợp lý.

• Các độ dịch chuyển anchor được tham số hoá như sau:

\[ t_x = \frac{x - x_a}{w_a}, \quad t_y = \frac{y - y_a}{h_a}, \quad t_w = \log\frac{w}{w_a}, \quad t_h = \log\frac{h}{h_a} \]

• trong đó \((x, y, w, h)\) là tâm và kích thước của khung dự đoán, và \((x_a, y_a, w_a, h_a)\) là anchor. Phép biến đổi log cho chiều rộng và chiều cao đảm bảo khung dự đoán luôn dương và làm hồi quy bất biến với tỷ lệ.

• Faster R-CNN huấn luyện với hàm mất mát đa nhiệm: mất mát phân loại (cross-entropy từ chương 05) cho nhãn lớp, cộng với mất mát L1 mượt (smooth L1 loss) cho hồi quy khung bao. Smooth L1 ít nhạy cảm với các điểm ngoại lai hơn L2:

\[ \text{smooth}_{L1}(x) = \begin{cases} 0.5x^2 & \text{nếu } |x| < 1 \\ |x| - 0.5 & \text{trường hợp còn lại} \end{cases} \]

• Mạng Kim Tự Tháp Đặc trưng (Feature Pyramid Networks - FPN) (Lin và cộng sự, 2017) giải quyết vấn đề đa tỷ lệ bằng cách xây dựng một đường dẫn từ trên xuống (top-down) với các kết nối ngang (lateral connections) hợp nhất ngữ nghĩa cấp cao với chi tiết không gian cấp thấp. Backbone tạo ra các feature map ở nhiều tỷ lệ (mỗi tầng pooling giảm một nửa độ phân giải). FPN thêm một đường dẫn từ trên xuống, mỗi mức nhận feature map được nội suy từ mức trên và hợp nhất với mức dưới lên tương ứng qua các tích chập 1x1 ngang. Kết quả là một kim tự tháp các feature map, mỗi mức có cả ngữ nghĩa mạnh và độ phân giải không gian tốt.

• Vật thể nhỏ được phát hiện từ các mức độ phân giải cao hơn của kim tự tháp; vật thể lớn từ các mức độ phân giải thấp hơn. FPN hiện là thành phần tiêu chuẩn trong hầu hết các kiến trúc phát hiện hiện đại.

• Bộ phát hiện một giai đoạn (one-stage detectors) bỏ qua bước đề xuất hoàn toàn, dự đoán nhãn lớp và khung bao trong một lượt duy nhất. Cách này nhanh hơn nhưng trước đây kém chính xác hơn bộ phát hiện hai giai đoạn, cho đến khi focal loss thu hẹp khoảng cách.

• YOLO (You Only Look Once, Redmon và cộng sự, 2016) chia ảnh thành một lưới \(S \times S\). Mỗi ô lưới dự đoán \(B\) khung bao và \(C\) xác suất lớp. Nếu tâm của một vật thể rơi vào một ô lưới, ô đó chịu trách nhiệm phát hiện nó. YOLO cực kỳ nhanh vì toàn bộ việc phát hiện chỉ là một lượt forward duy nhất không có giai đoạn đề xuất.

• YOLOv2 thêm anchor boxes, batch normalisation và huấn luyện đa tỷ lệ. YOLOv3 sử dụng Feature Pyramid Network và dự đoán ở ba tỷ lệ. YOLOv4-v8 tiếp tục cải thiện với backbone tốt hơn, mạng tổng hợp đường dẫn (path aggregation networks), và tăng cường dữ liệu mosaic (ghép bốn ảnh lại với nhau trong quá trình huấn luyện để tăng độ đa dạng ngữ cảnh).

• SSD (Single Shot MultiBox Detector, Liu và cộng sự, 2016) dự đoán ở nhiều tỷ lệ feature map trong backbone, sử dụng anchor box tại mỗi tỷ lệ. Các feature map sớm (độ phân giải cao) phát hiện vật thể nhỏ; các feature map muộn (độ phân giải thấp) phát hiện vật thể lớn. SSD nhanh hơn Faster R-CNN với độ chính xác cạnh tranh.

• RetinaNet (Lin và cộng sự, 2017) xác định vấn đề cốt lõi của bộ phát hiện một giai đoạn: mất cân bằng lớp (class imbalance). Đại đa số các anchor box tương ứng với nền, tạo ra các mẫu âm tính dễ (easy negatives) chiếm ưu thế trong mất mát và lấn át gradient từ các mẫu dương tính hiếm.

• Focal loss giải quyết vấn đề này bằng cách giảm trọng số của các mẫu dễ:

\[ \text{FL}(p_t) = -\alpha_t (1 - p_t)^\gamma \log(p_t) \]

• trong đó \(p_t\) là xác suất dự đoán cho lớp đúng. Khi mô hình tự tin và đúng (\(p_t\) cao), \((1 - p_t)^\gamma\) nhỏ, giảm đóng góp mất mát từ các mẫu âm tính dễ. Siêu tham số \(\gamma\) (thường là 2) kiểm soát mức độ giảm trọng số. Với \(\gamma = 0\), focal loss trở thành cross-entropy tiêu chuẩn. Với focal loss, RetinaNet đạt độ chính xác tương đương bộ phát hiện hai giai đoạn với tốc độ của bộ phát hiện một giai đoạn.

• Phát hiện không anchor (anchor-free detection) loại bỏ hoàn toàn anchor box, giảm việc điều chỉnh siêu tham số và đơn giản hoá đường ống.

• FCOS (Fully Convolutional One-Stage, Tian và cộng sự, 2019) dự đoán, tại mọi vị trí không gian của feature map, khoảng cách từ vị trí đó đến bốn cạnh của khung bao gần nhất (trái, trên, phải, dưới) cộng với nhãn lớp. Một điểm số centerness giảm trọng số các dự đoán xa tâm vật thể, cải thiện chất lượng. FCOS sử dụng FPN để xử lý đa tỷ lệ.

• CenterNet (Zhou và cộng sự, 2019) phát hiện vật thể dưới dạng điểm: nó dự đoán một bản đồ nhiệt nơi các đỉnh tương ứng với tâm vật thể, sau đó hồi quy chiều rộng và chiều cao tại mỗi đỉnh. Phát hiện trở thành ước lượng keypoint. Cách này thanh lịch và không cần anchor, nhưng đòi hỏi hậu xử lý bản đồ nhiệt cẩn thận.

• CornerNet phát hiện vật thể dưới dạng cặp góc (trên-trái và dưới-phải). Nó dự đoán hai bản đồ nhiệt (một cho mỗi loại góc) và sử dụng embedding kết hợp (associative embedding) để ghép các góc tương ứng thành khung bao. Điều này tránh cần anchor và xử lý vật thể có hình dạng bất kỳ.

• Phân đoạn ngữ nghĩa (semantic segmentation) gán một nhãn lớp cho mọi pixel trong ảnh. Không giống phát hiện (đầu ra khung bao), phân đoạn tạo ra một bản đồ pixel dày đặc. Một cảnh đường phố có thể gán nhãn mọi pixel là đường, vỉa hè, xe hơi, người đi bộ, toà nhà, bầu trời, v.v.

• Mạng Tích Chập Hoàn Toàn (Fully Convolutional Networks - FCN) (Long và cộng sự, 2015) thích ứng CNN phân loại cho phân đoạn bằng cách thay thế các tầng fully connected bằng các tầng tích chập, cho phép mạng đầu ra một bản đồ không gian thay vì một lớp duy nhất. Nội suy lên (upsampling) (qua tích chập chuyển vị hoặc nội suy song tuyến tính) khôi phục đầu ra về độ phân giải đầu vào. Các kết nối tắt (skip connections) từ các tầng trước đó thêm lại chi tiết không gian bị mất trong quá trình giảm mẫu.

• Tích chập chuyển vị (transposed convolution) (đôi khi gọi là "giải tích chập" - deconvolution) là đối tác nội suy lên của tích chập. Trong khi tích chập có stride giảm kích thước không gian, tích chập chuyển vị tăng chúng. Nó chèn các số 0 giữa các phần tử đầu vào và sau đó áp dụng một tích chập tiêu chuẩn, học cách nội suy lên một cách hiệu quả.

• U-Net (Ronneberger và cộng sự, 2015) giới thiệu một kiến trúc mã hoá-giải mã đối xứng với các kết nối tắt tại mọi cấp độ. Bộ mã hoá (đường dẫn co) giảm độ phân giải không gian trong khi tăng số kênh, giống hệt như CNN phân loại. Bộ giải mã (đường dẫn mở rộng) nội suy lên trở lại độ phân giải đầy đủ. Các kết nối tắt nối feature map bộ mã hoá với feature map bộ giải mã tại mỗi cấp độ, cung cấp chi tiết không gian tinh tế cho bộ giải mã. Sự kết hợp giữa ngữ nghĩa cấp cao và chi tiết cấp thấp này tạo ra các ranh giới phân đoạn sắc nét, chính xác.

• U-Net ban đầu được thiết kế cho phân đoạn ảnh y sinh (nơi dữ liệu huấn luyện khan hiếm) và kiến trúc của nó đã trở thành nền tảng cho nhiều mô hình sau này, bao gồm U-Net trong các mô hình khuếch tán tiềm ẩn (file 04).

• DeepLab (Chen và cộng sự, 2014-2018) giới thiệu hai cải tiến chính cho phân đoạn:

  • Tích chập rỗng (atrous/dilated convolution): tích chập tiêu chuẩn với các khoảng trống được chèn giữa các phần tử bộ lọc, được kiểm soát bởi tỷ lệ giãn nở \(r\). Một bộ lọc 3x3 với dilation \(r\) có trường tiếp nhận \((2r + 1) \times (2r + 1)\) trong khi chỉ dùng 9 tham số. Điều này nắm bắt ngữ cảnh ở nhiều tỷ lệ mà không cần giảm mẫu, bảo toàn độ phân giải không gian.

  • Atrous Spatial Pyramid Pooling (ASPP): áp dụng song song nhiều tích chập rỗng với các tỷ lệ giãn nở khác nhau (ví dụ: tỷ lệ 1, 6, 12, 18), nối các kết quả, và kết hợp với tích chập 1x1. ASPP nắm bắt ngữ cảnh ở nhiều tỷ lệ đồng thời, tương tự về tinh thần với khối Inception (file 02) nhưng sử dụng dilation thay vì các kích thước kernel khác nhau.

• DeepLab cũng sử dụng Trường Ngẫu nhiên Có điều kiện (Conditional Random Field - CRF) (chương 05) làm bước hậu xử lý để tinh chỉnh ranh giới phân đoạn bằng cách khuyến khích các pixel gần nhau về mặt không gian với màu sắc tương tự chia sẻ cùng nhãn.

• Phân đoạn thực thể (instance segmentation) kết hợp phát hiện và phân đoạn: nó xác định từng thực thể vật thể riêng lẻ và tạo ra một mặt nạ mức pixel cho mỗi thực thể. Hai xe hơi trong một cảnh nhận được hai mặt nạ riêng biệt, không chỉ "xe" cho cả hai.

• Mask R-CNN (He và cộng sự, 2017) mở rộng Faster R-CNN bằng cách thêm một đầu phân đoạn nhỏ dự đoán mặt nạ nhị phân cho mỗi vật thể được phát hiện. Kiến trúc là Faster R-CNN + một nhánh mask: nhánh mask lấy các đặc trưng RoI-pooled và đầu ra một mặt nạ nhị phân \(m \times m\) cho mỗi lớp. Nó sử dụng RoIAlign thay vì RoI pooling: nội suy song tuyến tính tại các điểm được lấy mẫu chính xác thay vì các ô lưới lượng tử hoá, tránh sự lệch không gian mà lượng tử hoá gây ra. Thay đổi nhỏ này cải thiện đáng kể chất lượng mặt nạ.

• Mask R-CNN được huấn luyện với mất mát đa nhiệm: mất mát phân loại + mất mát hồi quy khung bao + mất mát mặt nạ (binary cross-entropy trên từng pixel). Nhánh mask dự đoán một mặt nạ cho mỗi lớp một cách độc lập; chỉ mặt nạ tương ứng với lớp được dự đoán được sử dụng, tách biệt dự đoán mặt nạ khỏi phân loại và cải thiện cả hai.

• Phân đoạn Panoptic hợp nhất phân đoạn ngữ nghĩa và thực thể thành một tác vụ duy nhất. Mọi pixel nhận được cả nhãn lớp (ngữ nghĩa) và ID thực thể (instance, cho các lớp "vật" như xe hơi và người). Các lớp "chất liệu" (bầu trời, đường, cỏ) chỉ nhận nhãn ngữ nghĩa vì chúng là các vùng vô định hình không có các thực thể đếm được.

• Metric chất lượng panoptic (PQ) đánh giá điều này bằng cách phân rã thành chất lượng phân đoạn (IoU trung bình của các segment được ghép) và chất lượng nhận dạng (điểm F1 của các segment được ghép):

\[ \text{PQ} = \underbrace{\frac{\sum_{(p,g) \in \text{TP}} \text{IoU}(p,g)}{|\text{TP}|}}_{\text{SQ}} \times \underbrace{\frac{|\text{TP}|}{|\text{TP}| + \frac{1}{2}|\text{FP}| + \frac{1}{2}|\text{FN}|}}_{\text{RQ}} \]

• Phân đoạn thời gian thực rất quan trọng cho các ứng dụng như lái xe tự động và thực tế tăng cường, nơi ngân sách độ trễ rất chặt (thường dưới 30 mili giây mỗi khung hình).

• BiSeNet (Bilateral Segmentation Network, Yu và cộng sự, 2018) sử dụng hai đường dẫn song song: một đường dẫn không gian (spatial path) với các tầng rộng, nông giữ chi tiết không gian, và một đường dẫn ngữ cảnh (context path) với các tầng sâu, hẹp nắm bắt ngữ nghĩa. Các đầu ra được kết hợp, cho cả tốc độ và độ chính xác.

• DDRNet (Deep Dual-Resolution Network, Hong và cộng sự, 2021) duy trì hai nhánh ở các độ phân giải khác nhau xuyên suốt mạng, với trao đổi thông tin lặp đi lặp lại giữa chúng. Nhánh độ phân giải cao giữ chi tiết không gian trong khi nhánh độ phân giải thấp nắm bắt ngữ cảnh toàn cục. Nhiều mô-đun kết hợp song phương hai chiều hợp nhất thông tin theo cả hai hướng.

• Xu hướng chung trong phân đoạn thời gian thực là tránh mẫu mã hoá-giải mã nặng nề và thay vào đó duy trì đủ độ phân giải không gian xuyên suốt mạng, đánh đổi một phần độ chính xác để giảm đáng kể độ trễ.

Bài tập lập trình (dùng CoLab hoặc notebook)

1. Cài đặt tính toán IoU và Không có Cực đại Giả (NMS) từ đầu. Áp dụng NMS lên một tập các khung bao chồng lấn và trực quan hoá kết quả.

   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   import matplotlib.patches as patches
   
   def compute_iou(box1, box2):
       """Compute IoU between two boxes [x1, y1, x2, y2]."""
       x1 = jnp.maximum(box1[0], box2[0])
       y1 = jnp.maximum(box1[1], box2[1])
       x2 = jnp.minimum(box1[2], box2[2])
       y2 = jnp.minimum(box1[3], box2[3])
   
       intersection = jnp.maximum(0, x2 - x1) * jnp.maximum(0, y2 - y1)
       area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
       area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
       union = area1 + area2 - intersection
   
       return intersection / (union + 1e-6)
   
   def nms(boxes, scores, iou_threshold=0.5):
       """Non-Maximum Suppression."""
       order = jnp.argsort(-scores)  # sort by descending confidence
       keep = []
   
       remaining = list(range(len(scores)))
       order_list = order.tolist()
   
       while order_list:
           idx = order_list[0]
           keep.append(idx)
           order_list = order_list[1:]
   
           new_order = []
           for j in order_list:
               iou = compute_iou(boxes[idx], boxes[j])
               if iou < iou_threshold:
                   new_order.append(j)
           order_list = new_order
   
       return keep
   
   # Example: overlapping detections of the same object
   boxes = jnp.array([
       [50, 60, 150, 160],   # high confidence
       [55, 65, 155, 165],   # overlapping duplicate
       [52, 58, 148, 158],   # overlapping duplicate
       [200, 100, 300, 200], # different object
       [205, 105, 305, 205], # overlapping duplicate
   ])
   scores = jnp.array([0.95, 0.80, 0.70, 0.90, 0.60])
   
   keep = nms(boxes, scores, iou_threshold=0.5)
   
   fig, axes = plt.subplots(1, 2, figsize=(14, 5))
   colors = ['#3498db', '#e74c3c', '#27ae60', '#9b59b6', '#f39c12']
   
   for ax, title, indices in zip(axes, ['Before NMS', 'After NMS'],
                                  [range(len(boxes)), keep]):
       ax.set_xlim(0, 400); ax.set_ylim(0, 300)
       ax.set_aspect('equal'); ax.invert_yaxis()
       ax.set_title(title)
       for i in indices:
           b = boxes[i]
           rect = patches.Rectangle((b[0], b[1]), b[2]-b[0], b[3]-b[1],
                                     linewidth=2, edgecolor=colors[i],
                                     facecolor='none')
           ax.add_patch(rect)
           ax.text(b[0], b[1]-5, f'{scores[i]:.2f}', color=colors[i], fontsize=10)
   
   plt.tight_layout(); plt.show()
   print(f"Kept {len(keep)} of {len(boxes)} boxes after NMS")
   

2. Cài đặt Mạng Đề xuất Vùng (RPN) đơn giản. Cho một feature map, tạo anchor box ở nhiều tỷ lệ và tỷ lệ khung hình, dự đoán điểm số đối tượng và độ dịch chuyển khung.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   import matplotlib.patches as patches
   
   def generate_anchors(feature_h, feature_w, stride, scales, ratios):
       """Generate anchor boxes for each position on the feature map."""
       anchors = []
       for y in range(feature_h):
           for x in range(feature_w):
               cx = (x + 0.5) * stride
               cy = (y + 0.5) * stride
               for s in scales:
                   for r in ratios:
                       w = s * jnp.sqrt(r)
                       h = s / jnp.sqrt(r)
                       anchors.append([cx - w/2, cy - h/2, cx + w/2, cy + h/2])
       return jnp.array(anchors)
   
   def rpn_forward(feature_map, params):
       """Simplified RPN: predicts objectness and box offsets per anchor."""
       H, W, C = feature_map.shape
       n_anchors = params['cls_w'].shape[1]
   
       # Slide a 1x1 conv over the feature map (simplified)
       cls_scores = feature_map.reshape(-1, C) @ params['cls_w']  # (H*W, n_anchors)
       box_offsets = feature_map.reshape(-1, C) @ params['reg_w']  # (H*W, n_anchors*4)
   
       cls_scores = jax.nn.sigmoid(cls_scores)
       return cls_scores.ravel(), box_offsets.reshape(-1, 4)
   
   # Setup
   feature_h, feature_w, channels = 4, 4, 16
   stride = 16  # each feature map cell covers 16x16 pixels
   scales = [32, 64, 128]
   ratios = [0.5, 1.0, 2.0]
   n_anchors_per_pos = len(scales) * len(ratios)
   
   key = jax.random.PRNGKey(42)
   k1, k2, k3 = jax.random.split(key, 3)
   
   feature_map = jax.random.normal(k1, (feature_h, feature_w, channels))
   params = {
       'cls_w': jax.random.normal(k2, (channels, n_anchors_per_pos)) * 0.01,
       'reg_w': jax.random.normal(k3, (channels, n_anchors_per_pos * 4)) * 0.01,
   }
   
   anchors = generate_anchors(feature_h, feature_w, stride, scales, ratios)
   scores, offsets = rpn_forward(feature_map, params)
   
   print(f"Feature map: {feature_h}x{feature_w}, stride={stride}")
   print(f"Anchors per position: {n_anchors_per_pos}")
   print(f"Total anchors: {len(anchors)}")
   print(f"Objectness scores shape: {scores.shape}")
   print(f"Box offsets shape: {offsets.shape}")
   
   # Visualise anchors for one position
   fig, ax = plt.subplots(figsize=(6, 6))
   img_size = feature_h * stride
   ax.set_xlim(0, img_size); ax.set_ylim(0, img_size)
   ax.invert_yaxis(); ax.set_aspect('equal')
   
   pos_idx = feature_h // 2 * feature_w + feature_w // 2  # centre position
   colors = ['#3498db', '#e74c3c', '#27ae60']
   for i, s in enumerate(scales):
       for j, r in enumerate(ratios):
           idx = pos_idx * n_anchors_per_pos + i * len(ratios) + j
           a = anchors[idx]
           rect = patches.Rectangle((a[0], a[1]), a[2]-a[0], a[3]-a[1],
                                     linewidth=1.5, edgecolor=colors[i],
                                     facecolor='none', linestyle=['--', '-', ':'][j])
           ax.add_patch(rect)
   
   ax.scatter([img_size/2], [img_size/2], c='red', s=50, zorder=5)
   ax.set_title(f'Anchors at centre position\n3 scales × 3 ratios = {n_anchors_per_pos}')
   ax.grid(True, alpha=0.3)
   plt.tight_layout(); plt.show()
   

3. Cài đặt bộ mã hoá-giải mã U-Net đơn giản với các kết nối tắt cho phân đoạn 1D (gán nhãn nhị phân cho tín hiệu 1D).

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def conv1d_same(x, kernel):
       """1D convolution with same padding."""
       k = len(kernel)
       pad = k // 2
       x_pad = jnp.pad(x, pad, mode='edge')
       n = len(x)
       out = jnp.zeros(n)
       for i in range(n):
           out = out.at[i].set(jnp.sum(x_pad[i:i+k] * kernel))
       return out
   
   def downsample(x):
       return x[::2]
   
   def upsample(x, target_len):
       return jnp.interp(jnp.linspace(0, 1, target_len), jnp.linspace(0, 1, len(x)), x)
   
   def unet_1d(x, params):
       """Simplified 1D U-Net with 2 encoder/decoder levels."""
       # Encoder
       e1 = jnp.maximum(0, conv1d_same(x, params['enc1']))
       e1_down = downsample(e1)
   
       e2 = jnp.maximum(0, conv1d_same(e1_down, params['enc2']))
       e2_down = downsample(e2)
   
       # Bottleneck
       bottleneck = jnp.maximum(0, conv1d_same(e2_down, params['bottleneck']))
   
       # Decoder with skip connections
       d2_up = upsample(bottleneck, len(e2))
       d2 = jnp.maximum(0, conv1d_same(d2_up + e2, params['dec2']))  # skip connection
   
       d1_up = upsample(d2, len(e1))
       d1 = conv1d_same(d1_up + e1, params['dec1'])  # skip connection
   
       return jax.nn.sigmoid(d1)
   
   # Create signal with labelled regions
   n = 128
   t = jnp.linspace(0, 4 * jnp.pi, n)
   signal = jnp.sin(t) + 0.5 * jnp.sin(3 * t)
   labels = (signal > 0.5).astype(jnp.float32)  # binary segmentation target
   
   key = jax.random.PRNGKey(42)
   keys = jax.random.split(key, 5)
   params = {
       'enc1': jax.random.normal(keys[0], (5,)) * 0.3,
       'enc2': jax.random.normal(keys[1], (5,)) * 0.3,
       'bottleneck': jax.random.normal(keys[2], (3,)) * 0.3,
       'dec2': jax.random.normal(keys[3], (5,)) * 0.3,
       'dec1': jax.random.normal(keys[4], (5,)) * 0.3,
   }
   
   def loss_fn(params, signal, labels):
       pred = unet_1d(signal, params)
       return -jnp.mean(labels * jnp.log(pred + 1e-7) + (1 - labels) * jnp.log(1 - pred + 1e-7))
   
   grad_fn = jax.jit(jax.grad(loss_fn))
   lr = 0.05
   
   for step in range(500):
       grads = grad_fn(params, signal, labels)
       params = {k: params[k] - lr * grads[k] for k in params}
   
   pred = unet_1d(signal, params)
   
   fig, axes = plt.subplots(3, 1, figsize=(12, 7), sharex=True)
   axes[0].plot(t, signal, color='#3498db', linewidth=1.5)
   axes[0].set_title('Input Signal'); axes[0].set_ylabel('Value')
   
   axes[1].fill_between(t, 0, labels, alpha=0.3, color='#27ae60')
   axes[1].set_title('Ground Truth Labels'); axes[1].set_ylabel('Label')
   
   axes[2].plot(t, pred, color='#e74c3c', linewidth=1.5)
   axes[2].fill_between(t, 0, (pred > 0.5).astype(float), alpha=0.2, color='#e74c3c')
   axes[2].set_title('U-Net Prediction'); axes[2].set_ylabel('Probability')
   axes[2].set_xlabel('t')
   
   plt.tight_layout(); plt.show()
   print(f"Final loss: {loss_fn(params, signal, labels):.4f}")
   print(f"Pixel accuracy: {jnp.mean((pred > 0.5) == labels):.2%}")