Shinhaeng Lee

The Challenge

Detecting polyps in gastrointestinal endoscopy is a crucial but challenging task. Traditional detection methods often suffer from slow inference speed, high computational cost, and lower detection accuracy in real-time applications.

With the increasing demand for efficient and accurate AI-driven medical diagnostics, we needed a fast and lightweight solution to detect polyps without sacrificing accuracy.

The Solution: Pruned YOLOv3

Leveraging the power of YOLOv3, we designed an optimized deep learning model that improves real-time polyp detection by applying model pruning techniques.

🚀 Key Features

⚡ Model Pruning: Reduces redundant parameters for faster inference.
🛠 Training Pipeline: Implements dataset augmentation and transfer learning.
🎯 Real-Time Detection: Optimized for endoscopy applications.
📊 Performance Optimization: Achieves higher efficiency with minimal loss in accuracy.

🏗️ How It Works

We applied structured pruning to remove redundant convolutional filters, reducing the model size and computation while preserving detection accuracy.

📌 Step 1: Pruning the Model

# Load the pre-trained YOLOv3 model
model = load_yolov3_model("yolov3.weights")
 
# Apply structured pruning to convolutional layers
pruned_model = apply_pruning(model, prune_ratio=0.5)
 
# Save the optimized model
pruned_model.save("yolov3_pruned.pth")

Here, we remove 50% of the least important filters, significantly improving the inference speed.

📌 Step 2: Fine-Tuning on a Medical Dataset

After pruning, the model requires fine-tuning on polyp detection datasets to regain lost accuracy.

# Load dataset
dataset = load_polyp_dataset("polyp-detection")
 
# Train pruned YOLOv3 model
train_model(pruned_model, dataset, epochs=10, learning_rate=0.001)

This step ensures the model re-learns essential patterns without overfitting.

📌 Step 3: Real-Time Inference

Once optimized, the model can perform real-time detection with minimal latency.

# Load pruned model
model = load_model("yolov3_pruned.pth")
 
# Perform real-time polyp detection
detect_objects(model, input_video="endoscopy.mp4", confidence_threshold=0.5)

By integrating the pruned YOLOv3, we achieve a 2x faster inference speed with only 2% accuracy loss.

🎯 Future Improvements

Integrating Self-Attention: To improve detection accuracy by focusing on polyp regions.
Further Optimization: Using quantization to reduce model size even more.
Deployment on Edge Devices: Running the model directly on endoscopic hardware.

📥 Get Started

Clone the repository and start testing:

git clone https://github.com/OneOfCosmosMostWanted/YOLOv3_Prunning.git
cd YOLOv3_Prunning
pip install -r requirements.txt

To run the notebook:

jupyter notebook Github_V2_of_YOLOv3_Prunning.ipynb

💡 Want to contribute? Feel free to explore and optimize the model further!

🔗 GitHub Repository