xvanov/crater-model

Crater Detection Models

This repository contains trained models for NASA Crater Detection Challenge.

Current Best Score: 83.70% CDA (YOLO12x v2)

Models Overview

Model	Track	Purpose	Key Metric	Size
yolo12x_crater_v2_best.pt	-	Object detection	83.70% CDA ⭐	119 MB
yolo12m_best.pt	-	Object detection	82.30% CDA	50 MB
angle_refinement_best.pt	B	Angle prediction	5.02° MAE	45 MB
rim-unet-resnet34/	C	Rim segmentation	7.14° MAE, 86.23% CDA	98 MB
unet_best_traced.pt	-	Filled mask segmentation	-	27 MB
rtdetr-l-crater-100ep.pt	-	Transformer detection	66.9% CDA	130 MB

---

1. YOLO12x Crater Detector v2 (BEST) ⭐

Best Single-Model Detection - Use this for highest accuracy crater detection.

Metric	Value
Competition CDA	83.70%
mAP50	98.60%
mAP50-95	89.06%
Precision	96.80%
Recall	94.86%
Optimal Confidence	0.85

Training Details

• Base Model: yolo12x.pt (COCO pretrained)
• Training Data: 3,735 images
• Epochs: 150 (full training)
• Batch Size: 8
• Image Size: 640×640
• Parameters: ~60M

Key Findings

• +1.40% CDA improvement over YOLO12m
• Uniform confidence threshold (0.85) is optimal
• Size-dependent thresholds do NOT help
• Primary bottleneck: False positives (8.38 pts lost)
• Small crater recall limited (37.8%) due to resize

from ultralytics import YOLO
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(repo_id="xvanov/crater-model", filename="yolo12x_crater_v2_best.pt")
model = YOLO(model_path)
results = model.predict("image.png", conf=0.85)  # Use 0.85 for best CDA

---

2. YOLO12m Crater Detector

Smaller Detection Model - Use this for faster inference or limited resources.

Metric	Value
Competition CDA	82.30%
mAP50-95	0.694
mAP50	0.852
Precision	0.863
Recall	0.811

from ultralytics import YOLO
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(repo_id="xvanov/crater-model", filename="yolo12m_best.pt")
model = YOLO(model_path)
results = model.predict("image.png", conf=0.85)  # Use 0.85 for best CDA

---

3. Angle Refinement Network (Track B)

Best Angle Prediction - ResNet18-based network that predicts angle corrections.

Metric	Value
Overall MAE	5.02°
Elliptical MAE (a/b ≥ 1.3)	4.44°
Baseline MAE	27.84°
Improvement	6× better
CDA Improvement	+0.44% (validated on 425 images)

• Architecture: ResNet18 encoder + MLP head (unfrozen)
• Input: 128×128 grayscale ROI
• Output: cos(2θ), sin(2θ) for angle prediction
• Training: 45 epochs, 48,868 samples

import torch
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(repo_id="xvanov/crater-model", filename="angle_refinement_best.pt")
config_path = hf_hub_download(repo_id="xvanov/crater-model", filename="angle_refinement_config.json")

# Load with AngleRefinementNet class
checkpoint = torch.load(model_path, map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"])

---

4. Rim U-Net (Track C)

Rim Mask Segmentation - U-Net trained on annular rim masks for better angle extraction.

Metric	Value
Angle MAE	7.14°
IoU	0.81
CDA (with ellipticity gating)	86.23%

• Architecture: U-Net with ResNet34 encoder
• Input: 128×128 grayscale ROI
• Output: Binary rim mask
• Training: 150 epochs, ~156k samples

Key Innovation: Rim masks preserve ellipse orientation better than filled masks.

import torch
from huggingface_hub import hf_hub_download

# Download TorchScript model (recommended)
model_path = hf_hub_download(
    repo_id="xvanov/crater-model",
    filename="rim-unet-resnet34/best_traced.pt"
)
model = torch.jit.load(model_path)
model.eval()

Ellipticity Gating (Required)

For best results, use threshold 1.20:

• If mask aspect < 1.20: use bbox-derived angle (0° or 90°)
• If mask aspect >= 1.20: use mask-fitted angle

---

Track B vs Track C Comparison

Aspect	Track B (Angle Refinement)	Track C (Rim U-Net)
Angle MAE	5.02°	7.14°
CDA Impact	Pending integration	+0.6% (86.23%)
Size	45 MB	98 MB
Approach	Predicts angle delta	Predicts rim mask
When to use	Post-process bbox angles	Replace filled U-Net

Recommendation: Track B achieves better angle accuracy. Track C is useful if you need mask segmentation.

---

5. U-Net Crater Segmentation (Original)

Filled Mask Segmentation - Original U-Net for crater mask prediction.

• Architecture: U-Net with MobileNetV2 encoder
• Input: 128×128 grayscale ROI
• Output: Binary filled crater mask
• Format: TorchScript (traced)

import torch
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(repo_id="xvanov/crater-model", filename="unet_best_traced.pt")
model = torch.jit.load(model_path)

---

6. RT-DETR-L (Experimental)

Transformer Detector - Alternative to YOLO, lower performance.

Metric	Value
CDA	66.9%
mAP50-95	0.560
mAP50	0.707

---

Full Pipeline

Image (2592×2048)
    │
    ▼ Resize to 640×640
    │
┌─────────────────────────────────────┐
│ YOLO12x Detection (best accuracy)   │
│ → yolo12x_crater_v2_best.pt         │
│ → Threshold: 0.85 → 83.70% CDA      │
│ OR                                  │
│ YOLO12m Detection (faster)          │
│ → yolo12m_best.pt → 82.30% CDA      │
└─────────────────────────────────────┘
    │
    ▼ For each detection
    │
┌─────────────────────────────────────┐
│ Option A: BBox → Ellipse (Fastest)  │
│ → 83.70% CDA (YOLO12x)              │
├─────────────────────────────────────┤
│ Option B: + Rim U-Net + Gating      │
│ → 86.23% CDA                        │
├─────────────────────────────────────┤
│ Option C: + Angle Refinement        │
│ → Best angle accuracy (5.02° MAE)   │
└─────────────────────────────────────┘
    │
    ▼
Ellipse: (cx, cy, a, b, θ)

---

Competition Constraints

All models designed for:

• CPU-only inference (m8g.xlarge ARM)
• <5s/image inference time
• <6GB RAM usage
• Approved libraries only (torch, cv2, numpy, scipy)

---

Files

File	Description	Size
yolo12x_crater_v2_best.pt	YOLO12x v2 crater detector (BEST) ⭐	119 MB
yolo12m_best.pt	YOLO12m crater detector	50 MB
unet_best_traced.pt	U-Net filled mask (TorchScript)	27 MB
angle_refinement_best.pt	Track B angle refinement	45 MB
angle_refinement_config.json	Track B config	<1 KB
angle_refinement_evaluation.json	Track B metrics	<1 KB
rim-unet-resnet34/best_traced.pt	Track C rim U-Net (TorchScript)	98 MB
rim-unet-resnet34/pytorch_model.pt	Track C full checkpoint	294 MB
rim-unet-resnet34/README.md	Track C model card	5 KB
rtdetr-l-crater-100ep.pt	RT-DETR detector	130 MB
submission_yolo_best.pt	Original submission model	50 MB
yolo12x_crater_v2_epoch109_*.pt	YOLO12x intermediate checkpoint	119 MB

---

Citation

For research and competition use. Models developed for NASA Crater Detection Data Science Challenge.

---

Last updated: 2026-01-16