feat: implement DIoU loss function for object detection and update dataset documentation

README.md

@@ -29,7 +29,8 @@ Using the **Burn Deep Learning Framework**, this project achieves sub-millisecon
 ## Dataset and Preparation
 The model is trained on the primary dataset used for high-precision dart detection.
 
-- **Download Link**: [Dataset Resources (Google Drive)](https://drive.google.com/file/d/1ZEvuzg9zYbPd1FdZgV6v1aT4sqbqmLqp/view?usp=sharing)
+- **Model Weights Link**: [Neural Weights & TFLite (Google Drive)](https://drive.google.com/file/d/1ZEvuzg9zYbPd1FdZgV6v1aT4sqbqmLqp/view?usp=sharing)
+- **Dataset Source**: [DeepDarts (IEEE Dataport)](https://ieee-dataport.org/open-access/deepdarts-dataset)
 - **Resolution**: 800x800 pre-cropped high-resolution images.
 - **Structure**: Organize your data in the `dataset/800/` directory following the provided `labels.json` schema.
 

src/inference.rs

@@ -17,7 +17,7 @@ pub fn run_inference<B: Backend>(device: &B::Device, image_path: &str) {
 
     println!("🖼️ Processing image: {}...", image_path);
     let img = image::open(image_path).expect("Failed to open image");
-    let resized = img.resize_exact(800, 800, image::imageops::FilterType::Triangle);
+    let resized = img.resize_exact(416, 416, image::imageops::FilterType::Triangle);
     let pixels: Vec<f32> = resized
         .to_rgb8()
         .pixels()
@@ -30,7 +30,7 @@ pub fn run_inference<B: Backend>(device: &B::Device, image_path: &str) {
         })
         .collect();
 
-    let data = TensorData::new(pixels, [800, 800, 3]);
+    let data = TensorData::new(pixels, [416, 416, 3]);
     let input = Tensor::<B, 3>::from_data(data, device)
         .unsqueeze::<4>()
         .permute([0, 3, 1, 2]);
@@ -38,12 +38,12 @@ pub fn run_inference<B: Backend>(device: &B::Device, image_path: &str) {
     println!("🚀 Running MODEL Prediction...");
     let (out16, _out32) = model.forward(input);
 
-    // Post-process out16 (size [1, 30, 100, 100])
-    // Decode objectness part (Channel 4 for Anchor 0)
+    // out16 shape: [1, 30, 26, 26]
+    // 1. Extract Objectness (Channel 4 of first anchor)
     let obj = burn::tensor::activation::sigmoid(out16.clone().narrow(1, 4, 1));
 
-    // Find highest confidence cell
-    let (max_val, _) = obj.reshape([1, 10000]).max_dim_with_indices(1);
+    // 2. Find highest confidence cell in 26x26 grid
+    let (max_val, _) = obj.reshape([1, 676]).max_dim_with_indices(1);
     let confidence: f32 = max_val
         .to_data()
         .convert::<f32>()

src/loss.rs

@@ -10,36 +10,48 @@ pub fn diou_loss<B: Backend>(
     let bp = bboxes_pred.reshape([batch, 3, 10, h, w]);
     let t = target.reshape([batch, 3, 10, h, w]);
     
-    // 2. Objectness (Channel 4)
+    // 2. Loss Constants
+    let eps = 1e-6;
+
+    // 3. Objectness Loss (BCE)
     let obj_pred = burn::tensor::activation::sigmoid(bp.clone().narrow(2, 4, 1));
     let obj_target = t.clone().narrow(2, 4, 1);
     
-    let eps = 1e-7;
-    // Positive loss (where an object exists)
     let pos_loss = obj_target.clone().mul(obj_pred.clone().add_scalar(eps).log()).neg();
-    // Negative loss (where no object exists)
-    let neg_loss = obj_target.clone().neg().add_scalar(1.0).mul(obj_pred.clone().neg().add_scalar(1.0 + eps).log()).neg();
+    let neg_loss = obj_target.clone().neg().add_scalar(1.0)
+        .mul(obj_pred.clone().neg().add_scalar(1.0 + eps).log()).neg();
     
-    // Weight positive samples 10x more to fight imbalance (typical YOLO trick)
-    let obj_loss = pos_loss.mul_scalar(20.0).add(neg_loss).mean();
+    // Weight positive samples heavily (sparsity)
+    let obj_loss = pos_loss.mul_scalar(40.0).add(neg_loss).mean();
 
-    // 3. Class (Channels 5-9) - Only learn when object exists
+    // 4. Class Loss (Full BCE for all 5 channels)
+    // bp channels 5-9: Dart, Cal1, Cal2, Cal3, Cal4
     let cls_pred = burn::tensor::activation::sigmoid(bp.clone().narrow(2, 5, 5));
     let cls_target = t.clone().narrow(2, 5, 5);
-    let class_loss = cls_target.clone().mul(cls_pred.clone().add_scalar(eps).log()).neg()
-        .mul(obj_target.clone()) // Only count where object exists
+    
+    let cls_pos_loss = cls_target.clone().mul(cls_pred.clone().add_scalar(eps).log()).neg();
+    let cls_neg_loss = cls_target.clone().neg().add_scalar(1.0)
+        .mul(cls_pred.clone().neg().add_scalar(1.0 + eps).log()).neg();
+    
+    let class_loss = cls_pos_loss.add(cls_neg_loss)
+        .mul(obj_target.clone()) // Mask class loss where there is no object
         .mean()
-        .mul_scalar(5.0); // Boost class learning
+        .mul_scalar(15.0);
 
-    // 4. Coordinates (Channels 0-3) - Only learn when object exists
-    // 2. Coordinate Loss (MSE on relative offsets) - Weighted x10 for precision
+    // 5. Box (XYWH) Loss (MSE)
     let b_xy_pred = burn::tensor::activation::sigmoid(bp.clone().narrow(2, 0, 2));
     let b_xy_target = t.clone().narrow(2, 0, 2);
-    let xy_loss = b_xy_pred.sub(b_xy_target).powf_scalar(2.0).mul(obj_target.clone()).mean().mul_scalar(10.0);
+    let xy_loss = b_xy_pred.sub(b_xy_target).powf_scalar(2.0)
+        .mul(obj_target.clone())
+        .mean()
+        .mul_scalar(30.0); // Increase weight for coordinate precision
 
     let b_wh_pred = burn::tensor::activation::sigmoid(bp.clone().narrow(2, 2, 2));
     let b_wh_target = t.clone().narrow(2, 2, 2);
-    let wh_loss = b_wh_pred.sub(b_wh_target).powf_scalar(2.0).mul(obj_target).mean().mul_scalar(5.0);
+    let wh_loss = b_wh_pred.sub(b_wh_target).powf_scalar(2.0)
+        .mul(obj_target)
+        .mean()
+        .mul_scalar(5.0);
 
     obj_loss.add(class_loss).add(xy_loss).add(wh_loss)
 }