Update README.md

README.md

@@ -1,5 +1,5 @@
 ---
-license: apache-2.0
+license: mit
 language:
 - en
 base_model:
@@ -10,12 +10,157 @@ tags:
 - Threat_detection
 ---
 
-# Transformers for Object detection
+# RF-DETR based Threat Detection Model
 
-It's always been a truth that CNNs have faster inferencing than transformers, but this is now no more, Roboflow has release RF-Detr a transformer based object detection model, which not only outprforms CNNs, but also faster at accurate predictions.
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.0+-red.svg)](https://pytorch.org/)
+[![License](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
+[![RF-DETR](https://img.shields.io/badge/RF--DETR-Nano-orange.svg)](https://github.com/roboflow/rf-detr)
+[![mAP@50](https://img.shields.io/badge/mAP@50-84.8%25-brightgreen.svg)](#performance-metrics)
 
-This is the finetuned version of RF-Detr(Nano) model on the custom dataset for classification of Threat in out of four labels from: Gun(Including any kind of firearm weapon), Explosive(Fire or exploding scenario), Grenade(hand grenades), and Knife.
+## Transformers for Object Detection
 
+The paradigm has shifted! While CNNs traditionally dominated object detection with faster inference times, **RF-DETR** (Roboflow's Detection Transformer) has revolutionized the field. This transformer-based architecture not only **outperforms CNNs** in accuracy but also delivers **faster inference** for real-time applications.
+
+This repository contains a **fine-tuned RF-DETR Nano model** specifically trained for **threat detection**, capable of identifying four critical threat categories with high precision and speed.
+
+## Model Overview
+
+**RF-DETR Threat Detection** is a specialized computer vision model designed for security and surveillance applications. Built on Roboflow's cutting-edge RF-DETR architecture, this model can accurately detect and classify potential threats in real-time scenarios.
+
+The threat categories are as:
+
+| Class ID | Threat Type | Description |
+|----------|-------------|-------------|
+| 1 | **Gun** | Any type of firearm weapon including pistols, rifles, and other firearms |
+| 2 | **Explosive** | Fire, explosion scenarios, and explosive devices |
+| 3 | **Grenade** | Hand grenades and similar explosive devices |
+| 4 | **Knife** | Bladed weapons including knives, daggers, and sharp objects |
+
+## Training Dataset 
+
+Our custom threat detection dataset was meticulously curated and annotated to ensure robust model performance across diverse scenarios.
+
+### Class Distribution
+![Class Distribution](class_distribution.png)
+
+### Sample Annotations
+![Sample Annotations](sample_images_annotated.png)
+
+### Object Size Analysis
+![Object Size Distribution](Threat_COCO_dataset/visualizations/object_size_distribution.png)
+
+The model is trained to detect threats across various scales, from small concealed weapons to larger explosive devices.
+
+---
+## Performance Metrics
+
+### Training Performance
+![Training Metrics](metrics_plot.png)
+
+The training process demonstrates excellent convergence with:
+- **Consistent loss reduction** over 50 epochs
+- **Stable validation performance** indicating good generalization
+- **Balanced precision and recall** across all threat categories
+
+### Validation Results
+
+| Metric | Gun | Explosive | Grenade | Knife | **Overall** |
+|--------|-----|-----------|---------|-------|-------------|
+| **mAP@50:95** | 62.3% | 47.2% | 80.5% | 54.4% | **61.1%** |
+| **mAP@50** | 90.1% | 69.6% | 93.7% | 85.8% | **84.8%** |
+| **Precision** | 92.4% | 54.6% | 97.2% | 91.1% | **83.8%** |
+| **Recall** | 85.0% | 85.0% | 85.0% | 85.0% | **85.0%** |
+
+### Test Results
+
+| Metric | Gun | Explosive | Grenade | Knife | **Overall** |
+|--------|-----|-----------|---------|-------|-------------|
+| **mAP@50:95** | 65.3% | 35.7% | 83.2% | 49.8% | **58.5%** |
+| **mAP@50** | 93.1% | 60.5% | 91.1% | 79.7% | **81.1%** |
+| **Precision** | 96.7% | 49.7% | 93.1% | 86.5% | **81.5%** |
+| **Recall** | 83.0% | 83.0% | 83.0% | 83.0% | **83.0%** |
+
+### Key Performance Highlights
+
+- **84.8% mAP@50** on validation set
+- **Fast inference** with RF-DETR Nano architecture
+- **Excellent precision** for Gun (96.7%) and Grenade (93.1%) detection
+- **Consistent recall** of 83-85% across all threat categories
+- **Robust generalization** from validation to test performance
+
+
+
+# Install dependencies
+pip install torch torchvision
+pip install supervision
+pip install rfdetr
+pip install pillow requests numpy
+```
+
+### Basic Usage
+
+```python
+import numpy as np
+import supervision as sv
+from PIL import Image
+from rfdetr import RFDETRNano
+
+# Load the model
+model = RFDETRNano(
+    resolution=640, 
+    pretrain_weights="checkpoint_best_total.pth"
+)
+model.optimize_for_inference()
+
+# Load and process image
+image = Image.open("your_image.jpg")
+detections = model.predict(image, threshold=0.5)
+
+# Threat class mapping
+THREAT_CLASSES = {
+    1: "gun",
+    2: "explosive", 
+    3: "grenade",
+    4: "knife"
+}
+
+# Generate labels
+labels = [
+    f"{THREAT_CLASSES[class_id]} {confidence:.2f}"
+    for class_id, confidence in zip(detections.class_id, detections.confidence)
+]
+
+print(f"Detected {len(labels)} threats: {labels}")
+```
+
+
+## Model Architecture
+
+- **Base Architecture**: RF-DETR Nano
+- **Input Resolution**: 640×640 pixels
+- **Backbone**: Optimized transformer encoder
+- **Detection Head**: Custom 4-class threat detection
+- **Inference Speed**: ~50ms per image (GPU)
+- **Model Size**: Lightweight for edge deployment
+
+## Training Details
+
+### Training Configuration
+- **Epochs**: 50
+- **Batch Size**: Optimized for available GPU memory
+- **Optimizer**: AdamW with learning rate scheduling
+- **Data Augmentation**: Advanced augmentation pipeline for robust training
+- **Loss Function**: Multi-scale detection loss with class balancing
+
+### Training Strategy
+1. **Progressive Training**: Started with lower resolution, gradually increased
+2. **Class Balancing**: Weighted loss to handle class imbalance
+3. **Data Augmentation**: Extensive augmentation to improve generalization
+4. **Early Stopping**: Monitored validation mAP to prevent overfitting
+
+## Model Files
+
+- `checkpoint_best_total.pth` - Main model weights 
 
 ### Inference Instructions
 
@@ -86,4 +231,13 @@ annotated_image = bbox_annotator.annotate(annotated_image, detections)
 annotated_image = label_annotator.annotate(annotated_image, detections, labels)
 annotated_image.thumbnail((800, 800))
 annotated_image
-```
+```
+
+## Acknowledgments
+
+- **Roboflow** for the RF-DETR architecture
+- **Hugging Face** for model hosting and distribution
+- **PyTorch** ecosystem for deep learning framework
+- **Supervision** library for computer vision utilities
+
+**Disclaimer**: This model is designed for research purposes. It's predictions cannot be taken into account for deployment.