Duplicate from Subh775/Threat-Detection-RFDETR

Co-authored-by: Subhansh Malviya <Subh775@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+metrics_plot.png filter=lfs diff=lfs merge=lfs -text
+class_distribution.png filter=lfs diff=lfs merge=lfs -text
+sample_images_annotated.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,212 @@
+---
+license: mit
+language:
+- en
+base_model:
+- qualcomm/RF-DETR
+pipeline_tag: object-detection
+tags:
+- surveillance
+- Threat_detection
+---
+
+# RF-DETR based Threat Detection Model
+
+<a href="https://opensource.org/licenses/MIT">
+    <img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License">
+</a>
+<a href="https://github.com/roboflow/rf-detr">
+    <img src="https://img.shields.io/badge/RF--DETR-Nano-purple?logo=roboflow&logoColor=white" alt="Model">
+</a>
+<a href="#performance-metrics">
+    <img src="https://img.shields.io/badge/mAP%4050-84.8%25-darkgreen?style=flat" alt="mAP">
+</a>
+<a href="https://github.com/subh-775/Threat_Detection_YOLO-vs-RF-DETR">
+    <img src="https://img.shields.io/badge/-code-black?logo=github" alt="Code">
+</a>
+
+
+## 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.
+
+## Predicted Results
+![predictions](https://cdn-uploads.huggingface.co/production/uploads/66c6048d0bf40704e4159a23/MDRT7LUt1RQE60CGW8to4.jpeg)
+
+### Video Inferencing
+<video muted autoplay loop controls src="https://cdn-uploads.huggingface.co/production/uploads/66c6048d0bf40704e4159a23/5Kt3KghZaanzOVaVB6JS9.mp4" width=800></video>
+
+
+## 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](https://cdn-uploads.huggingface.co/production/uploads/66c6048d0bf40704e4159a23/5t7k-SJfuZWXJTek_RPWh.png)
+
+### Sample Annotations (Actual)
+![sample_images_annotated](https://cdn-uploads.huggingface.co/production/uploads/66c6048d0bf40704e4159a23/Mf65kxTEwfq9HPMlzwO5y.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
+
+## 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
+
+```python
+pip install -q rfdetr==1.2.1 supervision==0.26.1 
+```
+- You can use: [video_processing.py](https://huggingface.co/Subh775/Threat-Detection-RFDETR/blob/main/video_processing.py) to process large videos
+
+- Below is the script to process a single image
+
+```python
+import numpy as np
+import supervision as sv
+import torch
+import requests
+from PIL import Image
+import os
+
+from rfdetr import RFDETRNano
+
+THREAT_CLASSES = {
+    1: "Gun",
+    2: "Explosive", 
+    3: "Grenade",
+    4: "Knife"
+}
+
+image = Image.open("Path_to_image")
+
+# pre-trained weights
+weights_url = "https://huggingface.co/Subh775/Threat-Detection-RFDETR/resolve/main/checkpoint_best_total.pth"
+weights_filename = "checkpoint_best_total.pth"
+
+# Download weights if not already present
+if not os.path.exists(weights_filename):
+    print(f"Downloading weights from {weights_url}")
+    response = requests.get(weights_url, stream=True)
+    response.raise_for_status()
+    with open(weights_filename, 'wb') as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    print("Download complete.")
+
+model = RFDETRNano(resolution=640, pretrain_weights=weights_filename)
+model.optimize_for_inference()
+
+detections = model.predict(image, threshold=0.5)
+
+color = sv.ColorPalette.from_hex([
+    "#1E90FF", "#32CD32", "#FF0000", "#FF8C00"
+])
+
+text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
+thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)
+
+bbox_annotator = sv.BoxAnnotator(color=color, thickness=thickness)
+label_annotator = sv.LabelAnnotator(
+    color=color,
+    text_color=sv.Color.BLACK,
+    text_scale=text_scale,
+    smart_position=True
+)
+
+labels = []
+for class_id, confidence in zip(detections.class_id, detections.confidence):
+    class_name = THREAT_CLASSES.get(class_id, f"unknown_class_{class_id}")
+    labels.append(f"{class_name} {confidence:.2f}")
+
+annotated_image = image.copy()
+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 purpose only. It's predictions cannot be taken into account for deployment right now.

checkpoint.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15ce8511d2a3b5bb26c1e1cc8e16ecf11c6c3b6cc8e47fb4a7a56bb951d6c1da
+size 483371770

checkpoint0009.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4aab3f79814bdb5706a1307a68a3b3f78f9bf22afbe34b62e5d26b456ff8bf1
+size 483381074

checkpoint0019.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:729c30ee4cfedd0bd8ebc1d5d000abdbf8191f27522560b7777dd0aa9faf79da
+size 483381074

checkpoint0029.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b351c26d6c238fa5c77dfae8bd03379bfdfc0e68d98e420d6b433f4a7a4f03b
+size 483381074

checkpoint0039.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5d4fdbe25e9d7667040d62a52a38fed4797d1915e5e9df986ff80fc5d37802d
+size 483381074

checkpoint0049.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2494d99845cff1db236aa58be866d2748887ffefeaa17534ba48e25f820df6b3
+size 483381074

checkpoint_best_ema.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:955f076abe2011face8d84d34af1f003feec3455c22fa953250270291df082d7
+size 362564037

checkpoint_best_regular.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bafd5495a9a583d8a70a9ad66231bb1bf1665fc482f293dc56838cb01cad7fd
+size 362571545

checkpoint_best_total.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:481932df97ca0e3cbe73756a11a67dde6a617c548048baf67e29a9d640ed4a81
+size 120825206

eval/000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b734fccb5a7ff64331ceaae8d1bedfebf2ca7cf3c930868988168b883fb9a2d
+size 1196560

eval/latest.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fb5a3f509742430db35d4ca0c0c3afea92e9f33e7904b0eaf37f4b170f7c75b
+size 1193116

events.out.tfevents.1759683351.7100dacbcb24.36.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4392706209d595ace7553f208c3cecfeb915a9ae816726b8cedc30454680e8c7
+size 21772

results.json

@@ -0,0 +1 @@
+{"class_map": {"valid": [{"class": "Gun", "map@50:95": 0.622645348724507, "map@50": 0.9009196450890545, "precision": 0.9239130434782609, "recall": 0.85}, {"class": "explosion", "map@50:95": 0.4717619370767294, "map@50": 0.6955606171466157, "precision": 0.5463258785942492, "recall": 0.85}, {"class": "grenade", "map@50:95": 0.8045233282168012, "map@50": 0.9368179747971453, "precision": 0.9716312056737588, "recall": 0.85}, {"class": "knife", "map@50:95": 0.5443087767779382, "map@50": 0.8584450143594061, "precision": 0.9111111111111111, "recall": 0.85}, {"class": "all", "map@50:95": 0.6108098476989939, "map@50": 0.8479358128480553, "precision": 0.838245309714345, "recall": 0.85}], "test": [{"class": "Gun", "map@50:95": 0.6530858403662217, "map@50": 0.9310724801164318, "precision": 0.9672131147540983, "recall": 0.8300000000000001}, {"class": "explosion", "map@50:95": 0.3571121332596689, "map@50": 0.605062635831805, "precision": 0.4966887417218543, "recall": 0.8300000000000001}, {"class": "grenade", "map@50:95": 0.8318189072306177, "map@50": 0.9109914811178564, "precision": 0.9305555555555556, "recall": 0.8300000000000001}, {"class": "knife", "map@50:95": 0.49779614891880253, "map@50": 0.796572352660021, "precision": 0.8653846153846154, "recall": 0.8300000000000001}, {"class": "all", "map@50:95": 0.5849532574438278, "map@50": 0.8109247374315286, "precision": 0.814960506854031, "recall": 0.8300000000000001}]}, "map": 0.8479358128480553, "precision": 0.838245309714345, "recall": 0.85}

video_processing.py

@@ -0,0 +1,224 @@
+# pip install -q rfdetr==1.2.1 supervision==0.26.1
+
+# RF-DETR video processing for threat detection.
+# Inference time depends on frame resolution (e.g., ~50 ms/frame on GPU for 640×640).
+
+
+import numpy as np
+import supervision as sv
+import torch
+import requests
+from PIL import Image
+import os
+import cv2
+from tqdm import tqdm
+import time
+
+from rfdetr import RFDETRNano
+
+THREAT_CLASSES = {
+    1: "Gun",
+    2: "Explosive", 
+    3: "Grenade",
+    4: "Knife"
+}
+
+# Enable GPU if available
+if torch.cuda.is_available():
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
+    # print(f"CUDA Version: {torch.version.cuda}")
+    # print(f"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
+    
+    # Optimize for batch processing
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.deterministic = False
+else:
+    print("CUDA not available, using CPU")
+
+# Configuration
+INPUT_VIDEO = "test_video.mp4"
+
+base, ext = os.path.splitext(INPUT_VIDEO)
+OUTPUT_VIDEO = f"{base}_detr{ext}"
+
+THRESHOLD = 0.5
+BATCH_SIZE = 32  
+
+# Auto-adjust batch size based on GPU memory
+if torch.cuda.is_available():
+    gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
+
+print(f"Using batch size: {BATCH_SIZE}")
+
+# Download weights
+weights_url = "https://huggingface.co/Subh775/Threat-Detection-RFDETR/resolve/main/checkpoint_best_total.pth"
+weights_filename = "checkpoint_best_total.pth"
+
+if not os.path.exists(weights_filename):
+    print(f"Downloading weights from {weights_url}")
+    response = requests.get(weights_url, stream=True)
+    response.raise_for_status()
+    with open(weights_filename, 'wb') as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    print("Download complete.")
+
+print("Loading model...")
+model = RFDETRNano(resolution=640, pretrain_weights=weights_filename)
+model.optimize_for_inference()
+
+# Setup annotators
+color = sv.ColorPalette.from_hex([
+    "#1E90FF", "#32CD32", "#FF0000", "#FF8C00"
+])
+
+bbox_annotator = sv.BoxAnnotator(color=color, thickness=3)
+label_annotator = sv.LabelAnnotator(
+    color=color,
+    text_color=sv.Color.BLACK,
+    text_scale=1.0,
+    text_thickness=2,
+    smart_position=True
+)
+
+def process_frame_batch(frames):
+    """Process a batch of frames for better GPU utilization"""
+    batch_results = []
+    
+    # Convert all frames to PIL images
+    pil_images = []
+    for frame in frames:
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(rgb_frame)
+        pil_images.append(pil_image)
+    
+    # Process each image in the batch (RF-DETR processes them efficiently)
+    batch_detections = []
+    for pil_image in pil_images:
+        detections = model.predict(pil_image, threshold=THRESHOLD)
+        batch_detections.append(detections)
+    
+    # Annotate all images in the batch
+    annotated_frames = []
+    for pil_image, detections in zip(pil_images, batch_detections):
+        # Create labels
+        labels = []
+        for class_id, confidence in zip(detections.class_id, detections.confidence):
+            class_name = THREAT_CLASSES.get(class_id, f"unknown_class_{class_id}")
+            labels.append(f"{class_name} {confidence:.2f}")
+        
+        # Annotate
+        annotated_pil = pil_image.copy()
+        annotated_pil = bbox_annotator.annotate(annotated_pil, detections)
+        annotated_pil = label_annotator.annotate(annotated_pil, detections, labels)
+        
+        # Convert back to BGR
+        annotated_frame = cv2.cvtColor(np.array(annotated_pil), cv2.COLOR_RGB2BGR)
+        annotated_frames.append(annotated_frame)
+    
+    return annotated_frames, batch_detections
+
+# Open video
+cap = cv2.VideoCapture(INPUT_VIDEO)
+if not cap.isOpened():
+    print(f"Error: Could not open video file {INPUT_VIDEO}")
+    exit()
+
+# Get video properties
+fps = int(cap.get(cv2.CAP_PROP_FPS))
+width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+print(f"Video: {width}x{height}, {fps} FPS, {total_frames} frames")
+print(f"Processing in batches of {BATCH_SIZE} frames")
+
+# Setup video writer
+fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+out = cv2.VideoWriter(OUTPUT_VIDEO, fourcc, fps, (width, height))
+
+# Batch processing
+print("Processing video with batch inference...")
+frame_buffer = []
+total_detections = 0
+processed_frames = 0
+processing_times = []
+
+with tqdm(total=total_frames, desc="Batch processing") as pbar:
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            # Process remaining frames in buffer
+            if frame_buffer:
+                start_time = time.time()
+                annotated_frames, batch_detections = process_frame_batch(frame_buffer)
+                processing_time = time.time() - start_time
+                processing_times.append(processing_time)
+                
+                # Write remaining frames
+                for annotated_frame, detections in zip(annotated_frames, batch_detections):
+                    out.write(annotated_frame)
+                    total_detections += len(detections)
+                
+                processed_frames += len(frame_buffer)
+                pbar.update(len(frame_buffer))
+            break
+        
+        # Add frame to buffer
+        frame_buffer.append(frame)
+        
+        # Process when buffer is full
+        if len(frame_buffer) >= BATCH_SIZE:
+            start_time = time.time()
+            
+            # Process batch
+            annotated_frames, batch_detections = process_frame_batch(frame_buffer)
+            
+            processing_time = time.time() - start_time
+            processing_times.append(processing_time)
+            
+            # Write frames
+            batch_threats = 0
+            for annotated_frame, detections in zip(annotated_frames, batch_detections):
+                out.write(annotated_frame)
+                batch_threats += len(detections)
+                total_detections += len(detections)
+            
+            processed_frames += len(frame_buffer)
+            
+            # Update progress
+            batch_fps = len(frame_buffer) / processing_time if processing_time > 0 else 0
+            pbar.set_postfix({
+                'Batch FPS': f"{batch_fps:.1f}",
+                'Threats': batch_threats,
+                'Total': total_detections
+            })
+            pbar.update(len(frame_buffer))
+            
+            # Clear buffer
+            frame_buffer = []
+            
+            # Clear GPU cache every 10 batches
+            if torch.cuda.is_available() and processed_frames % (BATCH_SIZE * 10) == 0:
+                torch.cuda.empty_cache()
+
+# Cleanup
+cap.release()
+out.release()
+
+if torch.cuda.is_available():
+    torch.cuda.empty_cache()
+
+# Performance summary
+total_time = sum(processing_times)
+avg_fps = processed_frames / total_time if total_time > 0 else 0
+speedup = avg_fps / fps if fps > 0 else 0
+
+print(f"Output: {OUTPUT_VIDEO}")
+print(f"Stats:")
+print(f" • Processed: {processed_frames} frames")
+print(f" • Detections: {total_detections}")
+print(f" • Batch size: {BATCH_SIZE}")
+print(f" • Average speed: {avg_fps:.1f} FPS")
+print(f" • Speedup: {speedup:.1f}x real-time")
+print(f" • Processing time: {total_time:.1f}s")