DEBRIS DETECTION
[Baseline] Space Debris Detection
A getting started code for the Space Debris Detection Challenge
Shubhamaicrowd
Getting Started Code for Space Debris Challenge on AIcrowd¶
Author : Shubhamai¶
Download Necessary Packages 📚¶
In [ ]:
!pip install -U aicrowd-cli
# Installing TOrch
!pip install pyyaml==5.1
!pip install torch==1.7.1 torchvision==0.8.2
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
!gcc --version
# Installing Detectron2
import torch
assert torch.__version__.startswith("1.7")
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.7/index.html
Download Data¶
The first step is to download out train test data. We will be training a model on the train data and make predictions on test data. We submit our predictions.
Please enter your API Key from [https://www.aicrowd.com/participants/me]
In [ ]:
API_KEY = "22d25c05657ec939924bb37cc37ac813"
!aicrowd login --api-key $API_KEY
In [ ]:
!aicrowd dataset download --challenge debris-detection
In [ ]:
!rm -rf data
!mkdir data
!unzip train.zip -d data/train > /dev/null
!unzip val.zip -d data/val > /dev/null
!unzip test.zip -d data/test > /dev/null
!mv train.csv data/train.csv
!mv val.csv data/val.csv
Import packages¶
In [ ]:
import torch, torchvision
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
from detectron2 import model_zoo
from detectron2.engine import DefaultTrainer, DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.structures import BoxMode
import pandas as pd
import numpy as np
import os
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
import cv2
import random
from ast import literal_eval
In [ ]:
data_path = "data"
train_df = pd.read_csv(os.path.join(data_path, "train.csv"))
val_df = pd.read_csv(os.path.join(data_path, "val.csv"))
Visualize the data 👀¶
In [ ]:
train_df
Out[ ]:
In [ ]:
def show_images(images, num = 5):
images_to_show = np.random.choice(images, num)
for image_id in images_to_show:
image = Image.open(os.path.join(data_path, f"train/{image_id}.jpg"))
bboxes = literal_eval(train_df.loc[train_df['ImageID'] == image_id]['bboxes'].values[0])
draw = ImageDraw.Draw(image)
for bbox in bboxes:
draw.rectangle([bbox[0], bbox[2], bbox[1], bbox[3]], width=1)
plt.figure(figsize = (15,15))
plt.imshow(image)
plt.show()
show_images(train_df['ImageID'].unique(), num = 5)
Creating Dataset¶
In [ ]:
dict_dataset = []
def get_dataset_dics():
for index, row in train_df.iterrows():
image = Image.open(os.path.join(data_path, f"train/{row['ImageID']}.jpg"))
w, h = image.size
ann_lst = []
bboxes = literal_eval(row['bboxes'])
for n, bbox in enumerate(bboxes):
ann_dict = {'bbox': [bbox[0], bbox[2], bbox[1], bbox[3]],
'bbox_mode': BoxMode.XYXY_ABS,
'category_id': 0, #i[1]['category_id'].values[0],
'iscrowd': 0}
ann_lst.append(ann_dict)
image_dict = {'annotations': ann_lst,
'file_name': os.path.join(data_path, f"train/{row['ImageID']}.jpg"),
'height': h,
'image_id': row["ImageID"], #i[1]['image_category_id'].values[0],
'width': w}
dict_dataset.append(image_dict)
return dict_dataset
dict_dataset = get_dataset_dics()
In [ ]:
d = f"debris_train{np.random.randint(10000)}"
DatasetCatalog.register(d, lambda d=d : get_dataset_dics())
MetadataCatalog.get(d).set(thing_classes=["Debris"])
obj_metadata = MetadataCatalog.get(d)
In [ ]:
for i in random.sample(dict_dataset, 3):
img = cv2.imread(i["file_name"])
visualizer = Visualizer(img, metadata=obj_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(i)
plt.imshow(out.get_image())
Creating the Model¶
In [ ]:
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_DC5_3x.yaml"))
cfg.DATASETS.TRAIN = (d,)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_DC5_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 200
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
Train the Model¶
In [ ]:
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
In [ ]:
%load_ext tensorboard
%tensorboard --logdir output
Testing Phase 😅¶
We are almost done. We trained and validated on the training data. Now its the time to predict on test set and make a submission.# Prediction on Evaluation Set
Loading Pretrained Model¶
In [ ]:
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
predictor = DefaultPredictor(cfg)
Predict Test Set¶
Predict on the test set and you are all set to make the submission!
In [ ]:
test_imgs_paths = os.listdir(os.path.join(data_path, "test"))
predictions = {"ImageID":[], "bboxes":[]}
for test_img_path in tqdm(test_imgs_paths):
img = cv2.imread(os.path.join(data_path, "test", test_img_path))
h, w, _ = img.shape
model_predictions = predictor(img)
bboxes = model_predictions['instances'].pred_boxes.tensor.cpu().numpy().tolist()
scores = model_predictions['instances'].scores.cpu().numpy().tolist()
for n, bbox in enumerate(bboxes):
bboxes[n] == bbox.append(scores[n])
image_id = test_img_path.split('.')[0]
predictions['ImageID'].append(image_id)
predictions['bboxes'].append(bboxes)
Save the prediction to csv¶
In [ ]:
submission = pd.DataFrame(predictions)
submission
In [ ]:
submission.to_csv("submission.csv", index=False)
Making Direct Submission¶
In [ ]:
!aicrowd submission create -c debris-detection -f submission.csv
Content
Comments
You must login before you can post a comment.