Loading

Mask Prediction

Mask prediction with YOLOv5

AP of 0.96 in only 5 epochs

Zac

YOLOv5: https://github.com/ultralytics/yolov5

For improvement try:

  • Running on more epochs
  • Changing the model used (default here is large but can go bigger or smaller)
  • Hyperparameters in the masks.yaml file
  • Data augmentations

Download the files 💾

Download AIcrowd CLI

In [ ]:
!pip install aicrowd-cli
%load_ext aicrowd.magic

Login to AIcrowd

In [ ]:
%aicrowd login

Download data and install YOLOv5

In [ ]:
!rm -rf data
!mkdir data
%aicrowd ds dl -c mask-prediction -o data
In [4]:
!unzip -qq data/train.zip -d data
!unzip -qq data/val.zip -d data
!unzip -qq data/test.zip -d data
In [ ]:
!git clone https://github.com/ultralytics/yolov5  # clone
%cd yolov5
!pip install -r requirements.txt
In [6]:
import pandas as pd
import numpy as np
import os
In [7]:
name2id = {"surgical": 0, "N95": 1, "KN95": 2, "cloth": 3}
id2name = ["surgical", "N95",    "KN95", "cloth"]

Set up folder structure for YOLOv5

In [8]:
!mkdir "../data/train/images"
!mkdir "../data/train/labels"
!mkdir "../data/val/images"
!mkdir "../data/val/labels"
!mv ../data/train/*.jpg "../data/train/images"
!mv ../data/val/*.jpg "../data/val/images"
In [9]:
!echo 'train: ../data/train/images' > masks.yaml
!echo 'val: ../data/val/images' >> masks.yaml
!echo '' >> masks.yaml
!echo 'nc: 4' >> masks.yaml
!echo '' >> masks.yaml
!echo 'names: ["surgical", "N95", "KN95", "cloth"]' >> masks.yaml
In [21]:
for folder in ('train', 'val'):
    labels = pd.read_csv(f'../data/{folder}.csv')
    for name, bbox, masktype in zip(labels['ImageID'], labels['bbox'], labels['masktype']):
        bbox = [int(x)/512 for x in bbox[1:-1].split(",")]
        xmin, ymin, xmax, ymax = bbox
        bbox = [(xmax+xmin)/2, (ymax+ymin)/2, xmax-xmin, ymax-ymin]
        bbox = [str(x) for x in bbox]
        masktype = name2id[masktype]

        f = open(f"../data/{folder}/labels/{name}.txt", "w+")
        f.write(f"{masktype} {' '.join(bbox)}")
        f.close()
In [ ]:
#optional
%pip install -q wandb
import wandb
wandb.login()

Train the model

In [ ]:
!python train.py --img 512 --batch 25 --epochs 6 --data masks.yaml --weights yolov5l.pt --cache --optimizer=Adam

Make predictions

In [ ]:
!python detect.py --imgsz=512 --conf-thres=0 --iou-thres=0.5 --max-det=1 --save-txt --nosave --weights runs/train/exp/weights/best.pt --source ../data/test

Generating Prediction File

In [58]:
submission = pd.read_csv('../data/sample_submission.csv')
In [59]:
for i, row in enumerate(submission['ImageID']):
    x = open(f"runs/detect/exp/labels/{row}.txt").read().split()
    mtype = id2name[int(x[0])]
    x_center, y_center, width, height = [float(i)*512 for i in x[1:]]
    xmin, ymin, xmax, ymax = round(x_center - width/2), round(y_center - height/2), round(x_center + width/2), round(y_center + height/2)
    bbox = f"({xmin},{ymin},{xmax},{ymax})"
    submission.loc[i, 'bbox'] = bbox
    submission.loc[i, 'masktype'] = mtype
In [ ]:
submission.head()
In [61]:
!rm -rf assets
!mkdir assets
submission.to_csv(os.path.join("assets", "submission.csv"))

Submitting our Predictions

Note : Please save the notebook before submitting it (Ctrl + S)

In [62]:
%aicrowd notebook submit -c mask-prediction -a assets --no-verify
Using notebook: Mask Recognition Challenge for submission...
Removing existing files from submission directory...
Scrubbing API keys from the notebook...
Collecting notebook...


                                                  ╭─────────────────────────╮                                                   
                                                  │ Successfully submitted! │                                                   
                                                  ╰─────────────────────────╯                                                   
                                                        Important links                                                         
┌──────────────────┬───────────────────────────────────────────────────────────────────────────────────────────────────────────┐
│  This submission │ https://www.aicrowd.com/challenges/ai-blitz-xiii/problems/mask-prediction/submissions/174850              │
│                  │                                                                                                           │
│  All submissions │ https://www.aicrowd.com/challenges/ai-blitz-xiii/problems/mask-prediction/submissions?my_submissions=true │
│                  │                                                                                                           │
│      Leaderboard │ https://www.aicrowd.com/challenges/ai-blitz-xiii/problems/mask-prediction/leaderboards                    │
│                  │                                                                                                           │
│ Discussion forum │ https://discourse.aicrowd.com/c/ai-blitz-xiii                                                             │
│                  │                                                                                                           │
│   Challenge page │ https://www.aicrowd.com/challenges/ai-blitz-xiii/problems/mask-prediction                                 │
└──────────────────┴───────────────────────────────────────────────────────────────────────────────────────────────────────────┘
In [ ]:


Comments

You must login before you can post a comment.

Execute