Cassiterite Instance Segmentation

Instance segmentation model for detecting cassiterite minerals in microscopy images using Mask R-CNN with PyTorch.

🌐 Try Demo Online

Live Demo on HuggingFace 🚀

Try the model directly in your browser without any installation!

---

📖 Overview

This project implements instance segmentation for cassiterite mineral detection using Mask R-CNN. The model can identify and segment individual cassiterite particles in microscopy images.

Key Features:

• Multiple model variants (8 combinations)
• K-Fold cross-validation training
• Comprehensive evaluation with error analysis
• Support for tiled inference on large images
• Pre-trained on COCO, fine-tuned on cassiterite dataset

---

🏗️ Model Variants

The project includes 8 different model configurations, combining:

Component	Options
Architecture	Standard Mask R-CNN / Amodal Mask R-CNN
Backbone	ResNet50-FPN V1 / ResNet50-FPN V2
Optimizer	SGD / Adam

Model Types:

• Standard Mask R-CNN: Traditional instance segmentation approach
• Amodal Mask R-CNN: Enhanced with amodal branch for better occlusion handling using dilated convolutions

Example combinations:

• ResNet50-FPN-V1 + Amodal + SGD
• ResNet50-FPN-V2 + Standard + Adam
• (and 6 more variants)

---

🔧 Installation & Setup

Prerequisites

• Python 3.8 or higher
• CUDA-compatible GPU (recommended)

Option 1: Using Conda (Recommended)

# Create conda environment
conda create -n cassiterite python=3.10
conda activate cassiterite

# Install PyTorch with CUDA support
conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia

# Install other dependencies
pip install -r requirements.txt

Option 2: Using venv

# Create virtual environment
python -m venv venv

# Activate environment
# On Windows:
venv\Scripts\activate
# On Linux/Mac:
source venv/bin/activate

# Install dependencies
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

---

📂 Dataset Structure

Place your dataset in the following structure:

dataset142/
├── image1.jpg
├── image1.json
├── image2.jpg
├── image2.json
└── ...

Each annotation file should contain polygon coordinates for cassiterite instances in COCO format.

---

🚀 Training

Basic Training (K-Fold Cross-Validation)

python train.py \
  --data_dir dataset142 \
  --model_type amodal \
  --backbone resnet50_fpn_v2 \
  --optimizer adam \
  --lr 1e-4 \
  --epochs 50 \
  --n_folds 5

Training Different Model Variants

Amodal Model with SGD:

python train.py \
  --model_type amodal \
  --backbone resnet50_fpn_v2 \
  --optimizer sgd \
  --lr 0.01 \
  --checkpoint_dir checkpoints/v2_amodal_sgd

Standard Model with Adam:

python train.py \
  --model_type standard \
  --backbone resnet50_fpn_v1 \
  --optimizer adam \
  --lr 1e-4 \
  --checkpoint_dir checkpoints/v1_standard_adam

Key Parameters

• --model_type: standard or amodal
• --backbone: resnet50_fpn_v1 or resnet50_fpn_v2
• --optimizer: sgd or adam
• --lr: Learning rate (0.01 for SGD, 1e-4 for Adam recommended)
• --epochs: Number of training epochs (default: 50)
• --n_folds: Number of folds for cross-validation (default: 5)

---

📊 Evaluation

Standard Evaluation

Evaluate model on test images with detailed error analysis:

python test_evaluation.py \
  --checkpoint checkpoints/resnet50_fpn_v2_amodal_adam/fold1/best_model.pth \
  --model_type amodal \
  --backbone resnet50_fpn_v2 \
  --testset_dir testset \
  --output_dir test_evaluation/v2_amodal_adam

Tiled Evaluation (for Large Images)

For high-resolution images, use tiled inference to avoid memory issues:

python test_evaluation_tiled.py \
  --checkpoint checkpoints/resnet50_fpn_v2_amodal_adam/fold1/best_model.pth \
  --model_type amodal \
  --backbone resnet50_fpn_v2 \
  --testset_dir testset \
  --output_dir test_evaluation_tiled/v2_amodal_adam \
  --tile_size 512 \
  --overlap 64

Parameters:

• --tile_size: Size of each tile (default: 512)
• --overlap: Overlap between tiles in pixels (default: 64)

Evaluation Output

The evaluation generates:

• Metrics: mAP, IoU, Precision, Recall in JSON format
• Visualizations:
  • Detection results with bounding boxes and masks
  • Error maps showing classification of predictions
  • Confidence distribution plots

Error Classification:

• 🟢 Green: Correct detections
• 🟠 Orange: Wrong location (IoU < threshold)
• 🔴 Red: False positives
• 🔵 Blue: False negatives (missed objects)
• 🟡 Yellow: Low confidence detections

---

📈 Results Examples

Detection Results

Example of successful cassiterite detection and segmentation

Error Analysis

Error analysis visualization showing different types of prediction errors

Tiled Results

Tiled inference result on large image

Error analysis for tiled inference

---

📁 Project Structure

├── train.py                          # Training script with K-Fold CV
├── test_evaluation.py                # Standard evaluation script
├── test_evaluation_tiled.py          # Tiled inference for large images
├── model.py                          # Model architectures (Standard & Amodal)
├── datareaders.py                    # Dataset loading and augmentation
├── utils.py                          # Utility functions
├── requirements.txt                  # Python dependencies
├── test_evaluation_guide.md          # Detailed evaluation guide
├── test_evaluation_tiled_guide.md    # Tiled evaluation guide
├── figure/                           # Example result visualizations
└── checkpoints/                      # Saved model checkpoints

---

🎯 Quick Start Example

# 1. Clone/Download the repository
# 2. Install dependencies (see Installation section)

# 3. Train a model
python train.py \
  --model_type amodal \
  --backbone resnet50_fpn_v2 \
  --optimizer adam \
  --lr 1e-4 \
  --epochs 50

# 4. Evaluate the trained model
python test_evaluation.py \
  --checkpoint checkpoints/resnet50_fpn_v2_amodal_adam_lr1e-04/fold1/best_model.pth \
  --model_type amodal \
  --backbone resnet50_fpn_v2 \
  --testset_dir testset \
  --output_dir results

---

📝 Citation

If you use this code in your research, please cite:

@software{cassiterite_segmentation,
  title={Identification of Cassiterite Mineral in Microscope Images using the Mask R-CNN Model},
  author={Kevin Naufal Dany},
  year={2025},
  url={https://github.com/kevinnaufaldany/cassiterite-instance-segmentation}
}

---

📄 License

This project is available for academic and research purposes.

---

🔗 Links

• Demo: HuggingFace Space
• Detailed Guides: See kombinasi142.md and test_evaluation_tiled_guide.md

---

Happy Segmentation! 🎉