Gurveer05/moon-crater-boulder-detection-yolov5

YOLOv5 for Crater/Boulder detection on Moon

A yolov5s and a yolov5l model was trained on a labelled dataset of marked craters/boulders on moon. This was trained for the 3rd problem statement Automatic detection of craters & boulders from Orbiter High Resolution Camera(OHRC) images using AI/ML techniques of Bharatiya Antariksh Hackathon 2024. Only the finetuned yolov5l is present in this repo.

Sample Outputs

How to use

• Install yolov5:

pip install -U yolov5

• Load model and perform prediction:

import yolov5
# from PIL import Image

# load model
model = yolov5.load('Gurveer05/moon-crater-boulder-detection-yolov5')
  
# set model parameters
model.conf = 0.25  # NMS confidence threshold
model.iou = 0.45  # NMS IoU threshold
model.agnostic = False  # NMS class-agnostic
model.multi_label = False  # NMS multiple labels per box
model.max_det = 1000  # maximum number of detections per image

# set image
img = 'path/to/image' # or use: img = Image.open('/path/to/image')

# perform inference
results = model(img) # add size=640 if needed

# inference with test time augmentation
results = model(img, augment=True)

# parse results
predictions = results.pred[0]
boxes = predictions[:, :4] # x1, y1, x2, y2
scores = predictions[:, 4]
categories = predictions[:, 5]

# show detection bounding boxes on image
results.show()

# save results into "results/" folder
results.save(save_dir='results/')

• Finetune the model on your custom dataset:

yolov5 train --data data.yaml --img 640 --batch 16 --weights Gurveer05/moon-crater-boulder-detection-yolov5 --epochs 10

Miscellaneous

• Sample Working
• Dataset
• Code for training
• Output for OHRC images
• Sliced OHRC images input
• Code for inference on sliced OHRC images (Marked Images and Bounding Boxes CSV outputs)
• Code for inference on sliced OHRC images (lat/long .shp files for bounding boxes and center)