Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cleanup scripts and notebooks #118

Merged
merged 1 commit into from
Jun 14, 2021
Merged

Cleanup scripts and notebooks #118

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
285 changes: 196 additions & 89 deletions ...books/align_with_ultralytics_yolov5.ipynb → ...ow-to-align-with-ultralytics-yolov5.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,130 +2,166 @@
"cells": [
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
"import torch\n",
"\n",
"os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
"\n",
"device = torch.device('cpu')"
]
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": 2,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=14796495.0), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n"
]
}
],
"outputs": [],
"source": [
"import cv2\n",
"import torch\n",
"from yolort.utils.image_utils import letterbox, non_max_suppression\n",
"from yolort.utils import cv2_imshow, get_image_from_url, read_image_to_tensor,update_module_state_from_ultralytics\n",
"\n",
"\n",
"from yolort.utils.image_utils import letterbox, non_max_suppression\n",
"from yolort.utils import (\n",
" cv2_imshow,\n",
" get_image_from_url,\n",
" read_image_to_tensor,\n",
" update_module_state_from_ultralytics,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare image and model weights to test"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"is_half = False\n",
"device = torch.device('cpu')\n",
"img_raw = get_image_from_url(\"https://gitee.com/zhiqwang/yolov5-rt-stack/raw/master/test/assets/bus.jpg\")\n",
"ultralytics_weights_path = r\"weights\"\n",
"save_rt_weight_path = r'./weights/yolov5s_v40_updated.pt'\n",
"\n",
"if os.path.exists(ultralytics_weights_path) is not True:\n",
" os.makedirs(ultralytics_weights_path)\n",
"ultralytics_weights_path = f\"{ultralytics_weights_path}/yolov5s.pt\"\n",
"img = letterbox(img_raw, new_shape=(640, 640))[0]\n",
"img = read_image_to_tensor(img, is_half)\n",
"img = img.to(device)\n",
"# img_path = \"https://raw.githubusercontent.com/zhiqwang/yolov5-rt-stack/master/test/assets/bus.jpg\"\n",
"img_path = \"https://gitee.com/zhiqwang/yolov5-rt-stack/raw/master/test/assets/bus.jpg\" # For user in China\n",
"img_raw = get_image_from_url(img_path)\n",
"\n",
"#the weight's url in https://github.com/ultralytics/yolov5/releases\n",
"torch.hub.download_url_to_file(r'https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt',\n",
" ultralytics_weights_path,\n",
" progress=True)\n",
"yolort_weight_path = 'yolov5s_r40_updated.pt'\n",
"ultralytics_weights_path = \"yolov5s.pt\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can download the weight with following methods\n",
"\n",
"\n"
"```python\n",
"torch.hub.download_url_to_file(\n",
" r'https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt',\n",
" ultralytics_weights_path,\n",
" progress=True,\n",
")\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# config envrioment https://github.com/ultralytics/yolov5/issues/36\n",
"!pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt"
"# Preprocess\n",
"img = letterbox(img_raw, new_shape=(640, 640))[0]\n",
"img = read_image_to_tensor(img, is_half)\n",
"img = img.to(device)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load model as ultralytics and inference"
]
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Using cache found in /home/featurize/.cache/torch/hub/ultralytics_yolov5_master\n",
"Using cache found in /root/.cache/torch/hub/ultralytics_yolov5_master\n",
"Fusing layers... \n",
"Model Summary: 224 layers, 7266973 parameters, 0 gradients\n",
"Adding AutoShape... \n",
"YOLOv5 🚀 2021-6-14 torch 1.8.1+cu111 CPU\n",
"YOLOv5 🚀 v5.0-189-gdaab682 torch 1.8.1+cu102 CUDA:0 (Tesla P100-SXM2-16GB, 16276.25MB)\n",
"\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([[3.95028e+02, 2.28687e+02, 4.80437e+02, 5.25319e+02, 8.81427e-01, 0.00000e+00],\n",
" [1.31801e+02, 2.40823e+02, 2.05202e+02, 5.10122e+02, 8.74423e-01, 0.00000e+00],\n",
" [3.28450e+01, 2.39402e+02, 1.42193e+02, 5.31945e+02, 8.50409e-01, 0.00000e+00],\n",
" [1.81174e+01, 1.36144e+02, 4.74266e+02, 4.48792e+02, 7.12929e-01, 5.00000e+00],\n",
" [1.97863e-01, 2.94924e+02, 4.41639e+01, 5.27107e+02, 4.00531e-01, 0.00000e+00]])\n"
]
}
],
"source": [
"# 1.load ultiralytics model and inference\n",
"conf = 0.25\n",
"iou = 0.45\n",
"\n",
"model = torch.hub.load('ultralytics/yolov5', 'custom', path=ultralytics_weights_path)\n",
"model = model.to(device)\n",
"model.conf = conf # confidence threshold (0-1)\n",
"model.iou = iou # NMS IoU threshold (0-1)\n",
"model.classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for persons, cats and dogs\n",
"model.eval()\n",
"results = model(img.unsqueeze(dim=0))[0]\n",
"results = non_max_suppression(results, conf, iou, agnostic=True)[0]\n",
"print(results)\n",
"\n"
"\n",
"with torch.no_grad():\n",
" ultralytics_dets = model(img[None])[0]\n",
" ultralytics_dets = non_max_suppression(ultralytics_dets, conf, iou, agnostic=True)[0]"
]
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Detection results with ultralytics:\n",
"tensor([[3.95028e+02, 2.28687e+02, 4.80437e+02, 5.25319e+02, 8.81427e-01, 0.00000e+00],\n",
" [1.31801e+02, 2.40823e+02, 2.05202e+02, 5.10122e+02, 8.74422e-01, 0.00000e+00],\n",
" [3.28450e+01, 2.39402e+02, 1.42193e+02, 5.31945e+02, 8.50409e-01, 0.00000e+00],\n",
" [1.81174e+01, 1.36144e+02, 4.74266e+02, 4.48792e+02, 7.12929e-01, 5.00000e+00],\n",
" [1.97870e-01, 2.94924e+02, 4.41640e+01, 5.27107e+02, 4.00531e-01, 0.00000e+00]])\n"
]
}
],
"source": [
"print(f'Detection results with ultralytics:\\n{ultralytics_dets}')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Updating model weights from ultralytics to yolort"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Downloading: \"https://github.com/ultralytics/yolov5/archive/v4.0.zip\" to /home/featurize/.cache/torch/hub/v4.0.zip\n",
"Using cache found in /root/.cache/torch/hub/ultralytics_yolov5_v4.0\n",
"\n",
" from n params module arguments \n",
" 0 -1 1 3520 models.common.Focus [3, 32, 3] \n",
Expand Down Expand Up @@ -160,60 +196,131 @@
}
],
"source": [
"# 2.convert yolov5 from ultralytics to rtstack\n",
"model = update_module_state_from_ultralytics(arch='yolov5s',\n",
" version='v4.0',\n",
" custom_path_or_model=ultralytics_weights_path,\n",
" set_fp16=is_half)\n",
"torch.save(model.state_dict(),save_rt_weight_path)"
"\n",
"torch.save(model.state_dict(), yolort_weight_path)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load model as yolort and inference"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"from yolort.models.yolo import yolov5_darknet_pan_s_r40 as yolov5s\n",
"\n",
"model = yolov5s(score_thresh=conf, nms_thresh=iou)\n",
"model.load_state_dict(torch.load(yolort_weight_path))\n",
"\n",
"# Load model\n",
"if is_half:\n",
" model.half() # to FP16\n",
"model.eval()\n",
"\n",
"with torch.no_grad():\n",
" yolort_dets = model(img[None])"
]
},
{
"cell_type": "code",
"execution_count": 17,
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'scores': tensor([0.88143, 0.87442, 0.85041, 0.71293, 0.40053]), 'labels': tensor([0, 0, 0, 5, 0]), 'boxes': tensor([[3.95028e+02, 2.28687e+02, 4.80437e+02, 5.25319e+02],\n",
"Detection boxes with yolort:\n",
"tensor([[3.95028e+02, 2.28687e+02, 4.80437e+02, 5.25319e+02],\n",
" [1.31801e+02, 2.40823e+02, 2.05202e+02, 5.10122e+02],\n",
" [3.28450e+01, 2.39402e+02, 1.42193e+02, 5.31945e+02],\n",
" [1.81174e+01, 1.36144e+02, 4.74266e+02, 4.48792e+02],\n",
" [1.97876e-01, 2.94923e+02, 4.41640e+01, 5.27107e+02]])}]\n"
" [1.97876e-01, 2.94923e+02, 4.41640e+01, 5.27107e+02]])\n"
]
}
],
"source": [
"# 3.load rtstack model and inference\n",
"from yolort.models.yolo import yolov5_darknet_pan_s_r40 as yolov5s\n",
"model = yolov5s(pretrained=False, score_thresh=conf, nms_thresh=iou)\n",
"model.load_state_dict(torch.load(save_rt_weight_path))\n",
"\n",
"# Load model\n",
"if is_half:\n",
" model.half() # to FP16\n",
"model.eval()\n",
"with torch.no_grad():\n",
" model_out = model(img[None])\n",
"\n",
"print(model_out)"
"print(f\"Detection boxes with yolort:\\n{yolort_dets[0]['boxes']}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": []
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Detection scores with yolort:\n",
"tensor([0.88143, 0.87442, 0.85041, 0.71293, 0.40053])\n"
]
}
],
"source": [
"print(f\"Detection scores with yolort:\\n{yolort_dets[0]['scores']}\")"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Detection labels with yolort:\n",
"tensor([0, 0, 0, 5, 0])\n"
]
}
],
"source": [
"print(f\"Detection labels with yolort:\\n{yolort_dets[0]['labels']}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Varify the detection results between yolort and ultralytics"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Exported model has been tested, and the result looks good!\n"
]
}
],
"source": [
"torch.testing.assert_allclose(yolort_dets[0]['boxes'], ultralytics_dets[:, :4], rtol=1e-05, atol=1e-05)\n",
"\n",
"print(\"Exported model has been tested, and the result looks good!\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "base",
"display_name": "Python 3",
"language": "python",
"name": "base"
"name": "python3"
},
"language_info": {
"codemirror_mode": {
Expand All @@ -225,7 +332,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.4"
"version": "3.6.9"
}
},
"nbformat": 4,
Expand Down
Loading