YOLOv5 v5.0 release compatibility update for YOLOv3
This commit is contained in:
Glenn Jocher
parent
47ac6833ca
commit
4d0c2e6eee
@ -113,7 +113,7 @@ $ python detect.py --source data/images --weights yolov3.pt --conf 0.25
|
||||
|
||||
### PyTorch Hub
|
||||
|
||||
To run **batched inference** with YOLOv5 and [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36):
|
||||
To run **batched inference** with YOLOv3 and [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36):
|
||||
```python
|
||||
import torch
|
||||
|
||||
|
||||
55
data/GlobalWheat2020.yaml
Normal file
55
data/GlobalWheat2020.yaml
Normal file
@ -0,0 +1,55 @@
|
||||
# Global Wheat 2020 dataset http://www.global-wheat.com/
|
||||
# Train command: python train.py --data GlobalWheat2020.yaml
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /datasets/GlobalWheat2020
|
||||
# /yolov3
|
||||
|
||||
|
||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
||||
train: # 3422 images
|
||||
- ../datasets/GlobalWheat2020/images/arvalis_1
|
||||
- ../datasets/GlobalWheat2020/images/arvalis_2
|
||||
- ../datasets/GlobalWheat2020/images/arvalis_3
|
||||
- ../datasets/GlobalWheat2020/images/ethz_1
|
||||
- ../datasets/GlobalWheat2020/images/rres_1
|
||||
- ../datasets/GlobalWheat2020/images/inrae_1
|
||||
- ../datasets/GlobalWheat2020/images/usask_1
|
||||
|
||||
val: # 748 images (WARNING: train set contains ethz_1)
|
||||
- ../datasets/GlobalWheat2020/images/ethz_1
|
||||
|
||||
test: # 1276 images
|
||||
- ../datasets/GlobalWheat2020/images/utokyo_1
|
||||
- ../datasets/GlobalWheat2020/images/utokyo_2
|
||||
- ../datasets/GlobalWheat2020/images/nau_1
|
||||
- ../datasets/GlobalWheat2020/images/uq_1
|
||||
|
||||
# number of classes
|
||||
nc: 1
|
||||
|
||||
# class names
|
||||
names: [ 'wheat_head' ]
|
||||
|
||||
|
||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
||||
download: |
|
||||
from utils.general import download, Path
|
||||
|
||||
# Download
|
||||
dir = Path('../datasets/GlobalWheat2020') # dataset directory
|
||||
urls = ['https://zenodo.org/record/4298502/files/global-wheat-codalab-official.zip',
|
||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/GlobalWheat2020_labels.zip']
|
||||
download(urls, dir=dir)
|
||||
|
||||
# Make Directories
|
||||
for p in 'annotations', 'images', 'labels':
|
||||
(dir / p).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Move
|
||||
for p in 'arvalis_1', 'arvalis_2', 'arvalis_3', 'ethz_1', 'rres_1', 'inrae_1', 'usask_1', \
|
||||
'utokyo_1', 'utokyo_2', 'nau_1', 'uq_1':
|
||||
(dir / p).rename(dir / 'images' / p) # move to /images
|
||||
f = (dir / p).with_suffix('.json') # json file
|
||||
if f.exists():
|
||||
f.rename((dir / 'annotations' / p).with_suffix('.json')) # move to /annotations
|
||||
52
data/SKU-110K.yaml
Normal file
52
data/SKU-110K.yaml
Normal file
@ -0,0 +1,52 @@
|
||||
# SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19
|
||||
# Train command: python train.py --data SKU-110K.yaml
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /datasets/SKU-110K
|
||||
# /yolov3
|
||||
|
||||
|
||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
||||
train: ../datasets/SKU-110K/train.txt # 8219 images
|
||||
val: ../datasets/SKU-110K/val.txt # 588 images
|
||||
test: ../datasets/SKU-110K/test.txt # 2936 images
|
||||
|
||||
# number of classes
|
||||
nc: 1
|
||||
|
||||
# class names
|
||||
names: [ 'object' ]
|
||||
|
||||
|
||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
||||
download: |
|
||||
import shutil
|
||||
from tqdm import tqdm
|
||||
from utils.general import np, pd, Path, download, xyxy2xywh
|
||||
|
||||
# Download
|
||||
datasets = Path('../datasets') # download directory
|
||||
urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz']
|
||||
download(urls, dir=datasets, delete=False)
|
||||
|
||||
# Rename directories
|
||||
dir = (datasets / 'SKU-110K')
|
||||
if dir.exists():
|
||||
shutil.rmtree(dir)
|
||||
(datasets / 'SKU110K_fixed').rename(dir) # rename dir
|
||||
(dir / 'labels').mkdir(parents=True, exist_ok=True) # create labels dir
|
||||
|
||||
# Convert labels
|
||||
names = 'image', 'x1', 'y1', 'x2', 'y2', 'class', 'image_width', 'image_height' # column names
|
||||
for d in 'annotations_train.csv', 'annotations_val.csv', 'annotations_test.csv':
|
||||
x = pd.read_csv(dir / 'annotations' / d, names=names).values # annotations
|
||||
images, unique_images = x[:, 0], np.unique(x[:, 0])
|
||||
with open((dir / d).with_suffix('.txt').__str__().replace('annotations_', ''), 'w') as f:
|
||||
f.writelines(f'./images/{s}\n' for s in unique_images)
|
||||
for im in tqdm(unique_images, desc=f'Converting {dir / d}'):
|
||||
cls = 0 # single-class dataset
|
||||
with open((dir / 'labels' / im).with_suffix('.txt'), 'a') as f:
|
||||
for r in x[images == im]:
|
||||
w, h = r[6], r[7] # image width, height
|
||||
xywh = xyxy2xywh(np.array([[r[1] / w, r[2] / h, r[3] / w, r[4] / h]]))[0] # instance
|
||||
f.write(f"{cls} {xywh[0]:.5f} {xywh[1]:.5f} {xywh[2]:.5f} {xywh[3]:.5f}\n") # write label
|
||||
61
data/VisDrone.yaml
Normal file
61
data/VisDrone.yaml
Normal file
@ -0,0 +1,61 @@
|
||||
# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset
|
||||
# Train command: python train.py --data VisDrone.yaml
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /VisDrone
|
||||
# /yolov3
|
||||
|
||||
|
||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
||||
train: ../VisDrone/VisDrone2019-DET-train/images # 6471 images
|
||||
val: ../VisDrone/VisDrone2019-DET-val/images # 548 images
|
||||
test: ../VisDrone/VisDrone2019-DET-test-dev/images # 1610 images
|
||||
|
||||
# number of classes
|
||||
nc: 10
|
||||
|
||||
# class names
|
||||
names: [ 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor' ]
|
||||
|
||||
|
||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
||||
download: |
|
||||
from utils.general import download, os, Path
|
||||
|
||||
def visdrone2yolo(dir):
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
|
||||
def convert_box(size, box):
|
||||
# Convert VisDrone box to YOLO xywh box
|
||||
dw = 1. / size[0]
|
||||
dh = 1. / size[1]
|
||||
return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh
|
||||
|
||||
(dir / 'labels').mkdir(parents=True, exist_ok=True) # make labels directory
|
||||
pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')
|
||||
for f in pbar:
|
||||
img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
|
||||
lines = []
|
||||
with open(f, 'r') as file: # read annotation.txt
|
||||
for row in [x.split(',') for x in file.read().strip().splitlines()]:
|
||||
if row[4] == '0': # VisDrone 'ignored regions' class 0
|
||||
continue
|
||||
cls = int(row[5]) - 1
|
||||
box = convert_box(img_size, tuple(map(int, row[:4])))
|
||||
lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
|
||||
with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl:
|
||||
fl.writelines(lines) # write label.txt
|
||||
|
||||
|
||||
# Download
|
||||
dir = Path('../VisDrone') # dataset directory
|
||||
urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip',
|
||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip',
|
||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip',
|
||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip']
|
||||
download(urls, dir=dir)
|
||||
|
||||
# Convert
|
||||
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
|
||||
visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels
|
||||
@ -1,9 +1,9 @@
|
||||
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
|
||||
# Train command: python train.py --data argoverse_hd.yaml
|
||||
# Default dataset location is next to /yolov5:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /argoverse
|
||||
# /yolov5
|
||||
# /yolov3
|
||||
|
||||
|
||||
# download command/URL (optional)
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
# COCO 2017 dataset http://cocodataset.org
|
||||
# Train command: python train.py --data coco.yaml
|
||||
# Default dataset location is next to /yolov3:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /coco
|
||||
# /yolov3
|
||||
@ -30,6 +30,6 @@ names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', '
|
||||
|
||||
# Print classes
|
||||
# with open('data/coco.yaml') as f:
|
||||
# d = yaml.load(f, Loader=yaml.FullLoader) # dict
|
||||
# d = yaml.safe_load(f) # dict
|
||||
# for i, x in enumerate(d['names']):
|
||||
# print(i, x)
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
# COCO 2017 dataset http://cocodataset.org - first 128 training images
|
||||
# Train command: python train.py --data coco128.yaml
|
||||
# Default dataset location is next to /yolov3:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /coco128
|
||||
# /yolov3
|
||||
|
||||
28
data/hyp.finetune_objects365.yaml
Normal file
28
data/hyp.finetune_objects365.yaml
Normal file
@ -0,0 +1,28 @@
|
||||
lr0: 0.00258
|
||||
lrf: 0.17
|
||||
momentum: 0.779
|
||||
weight_decay: 0.00058
|
||||
warmup_epochs: 1.33
|
||||
warmup_momentum: 0.86
|
||||
warmup_bias_lr: 0.0711
|
||||
box: 0.0539
|
||||
cls: 0.299
|
||||
cls_pw: 0.825
|
||||
obj: 0.632
|
||||
obj_pw: 1.0
|
||||
iou_t: 0.2
|
||||
anchor_t: 3.44
|
||||
anchors: 3.2
|
||||
fl_gamma: 0.0
|
||||
hsv_h: 0.0188
|
||||
hsv_s: 0.704
|
||||
hsv_v: 0.36
|
||||
degrees: 0.0
|
||||
translate: 0.0902
|
||||
scale: 0.491
|
||||
shear: 0.0
|
||||
perspective: 0.0
|
||||
flipud: 0.0
|
||||
fliplr: 0.5
|
||||
mosaic: 1.0
|
||||
mixup: 0.0
|
||||
102
data/objects365.yaml
Normal file
102
data/objects365.yaml
Normal file
@ -0,0 +1,102 @@
|
||||
# Objects365 dataset https://www.objects365.org/
|
||||
# Train command: python train.py --data objects365.yaml
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /datasets/objects365
|
||||
# /yolov3
|
||||
|
||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
||||
train: ../datasets/objects365/images/train # 1742289 images
|
||||
val: ../datasets/objects365/images/val # 5570 images
|
||||
|
||||
# number of classes
|
||||
nc: 365
|
||||
|
||||
# class names
|
||||
names: [ 'Person', 'Sneakers', 'Chair', 'Other Shoes', 'Hat', 'Car', 'Lamp', 'Glasses', 'Bottle', 'Desk', 'Cup',
|
||||
'Street Lights', 'Cabinet/shelf', 'Handbag/Satchel', 'Bracelet', 'Plate', 'Picture/Frame', 'Helmet', 'Book',
|
||||
'Gloves', 'Storage box', 'Boat', 'Leather Shoes', 'Flower', 'Bench', 'Potted Plant', 'Bowl/Basin', 'Flag',
|
||||
'Pillow', 'Boots', 'Vase', 'Microphone', 'Necklace', 'Ring', 'SUV', 'Wine Glass', 'Belt', 'Monitor/TV',
|
||||
'Backpack', 'Umbrella', 'Traffic Light', 'Speaker', 'Watch', 'Tie', 'Trash bin Can', 'Slippers', 'Bicycle',
|
||||
'Stool', 'Barrel/bucket', 'Van', 'Couch', 'Sandals', 'Basket', 'Drum', 'Pen/Pencil', 'Bus', 'Wild Bird',
|
||||
'High Heels', 'Motorcycle', 'Guitar', 'Carpet', 'Cell Phone', 'Bread', 'Camera', 'Canned', 'Truck',
|
||||
'Traffic cone', 'Cymbal', 'Lifesaver', 'Towel', 'Stuffed Toy', 'Candle', 'Sailboat', 'Laptop', 'Awning',
|
||||
'Bed', 'Faucet', 'Tent', 'Horse', 'Mirror', 'Power outlet', 'Sink', 'Apple', 'Air Conditioner', 'Knife',
|
||||
'Hockey Stick', 'Paddle', 'Pickup Truck', 'Fork', 'Traffic Sign', 'Balloon', 'Tripod', 'Dog', 'Spoon', 'Clock',
|
||||
'Pot', 'Cow', 'Cake', 'Dinning Table', 'Sheep', 'Hanger', 'Blackboard/Whiteboard', 'Napkin', 'Other Fish',
|
||||
'Orange/Tangerine', 'Toiletry', 'Keyboard', 'Tomato', 'Lantern', 'Machinery Vehicle', 'Fan',
|
||||
'Green Vegetables', 'Banana', 'Baseball Glove', 'Airplane', 'Mouse', 'Train', 'Pumpkin', 'Soccer', 'Skiboard',
|
||||
'Luggage', 'Nightstand', 'Tea pot', 'Telephone', 'Trolley', 'Head Phone', 'Sports Car', 'Stop Sign',
|
||||
'Dessert', 'Scooter', 'Stroller', 'Crane', 'Remote', 'Refrigerator', 'Oven', 'Lemon', 'Duck', 'Baseball Bat',
|
||||
'Surveillance Camera', 'Cat', 'Jug', 'Broccoli', 'Piano', 'Pizza', 'Elephant', 'Skateboard', 'Surfboard',
|
||||
'Gun', 'Skating and Skiing shoes', 'Gas stove', 'Donut', 'Bow Tie', 'Carrot', 'Toilet', 'Kite', 'Strawberry',
|
||||
'Other Balls', 'Shovel', 'Pepper', 'Computer Box', 'Toilet Paper', 'Cleaning Products', 'Chopsticks',
|
||||
'Microwave', 'Pigeon', 'Baseball', 'Cutting/chopping Board', 'Coffee Table', 'Side Table', 'Scissors',
|
||||
'Marker', 'Pie', 'Ladder', 'Snowboard', 'Cookies', 'Radiator', 'Fire Hydrant', 'Basketball', 'Zebra', 'Grape',
|
||||
'Giraffe', 'Potato', 'Sausage', 'Tricycle', 'Violin', 'Egg', 'Fire Extinguisher', 'Candy', 'Fire Truck',
|
||||
'Billiards', 'Converter', 'Bathtub', 'Wheelchair', 'Golf Club', 'Briefcase', 'Cucumber', 'Cigar/Cigarette',
|
||||
'Paint Brush', 'Pear', 'Heavy Truck', 'Hamburger', 'Extractor', 'Extension Cord', 'Tong', 'Tennis Racket',
|
||||
'Folder', 'American Football', 'earphone', 'Mask', 'Kettle', 'Tennis', 'Ship', 'Swing', 'Coffee Machine',
|
||||
'Slide', 'Carriage', 'Onion', 'Green beans', 'Projector', 'Frisbee', 'Washing Machine/Drying Machine',
|
||||
'Chicken', 'Printer', 'Watermelon', 'Saxophone', 'Tissue', 'Toothbrush', 'Ice cream', 'Hot-air balloon',
|
||||
'Cello', 'French Fries', 'Scale', 'Trophy', 'Cabbage', 'Hot dog', 'Blender', 'Peach', 'Rice', 'Wallet/Purse',
|
||||
'Volleyball', 'Deer', 'Goose', 'Tape', 'Tablet', 'Cosmetics', 'Trumpet', 'Pineapple', 'Golf Ball',
|
||||
'Ambulance', 'Parking meter', 'Mango', 'Key', 'Hurdle', 'Fishing Rod', 'Medal', 'Flute', 'Brush', 'Penguin',
|
||||
'Megaphone', 'Corn', 'Lettuce', 'Garlic', 'Swan', 'Helicopter', 'Green Onion', 'Sandwich', 'Nuts',
|
||||
'Speed Limit Sign', 'Induction Cooker', 'Broom', 'Trombone', 'Plum', 'Rickshaw', 'Goldfish', 'Kiwi fruit',
|
||||
'Router/modem', 'Poker Card', 'Toaster', 'Shrimp', 'Sushi', 'Cheese', 'Notepaper', 'Cherry', 'Pliers', 'CD',
|
||||
'Pasta', 'Hammer', 'Cue', 'Avocado', 'Hamimelon', 'Flask', 'Mushroom', 'Screwdriver', 'Soap', 'Recorder',
|
||||
'Bear', 'Eggplant', 'Board Eraser', 'Coconut', 'Tape Measure/Ruler', 'Pig', 'Showerhead', 'Globe', 'Chips',
|
||||
'Steak', 'Crosswalk Sign', 'Stapler', 'Camel', 'Formula 1', 'Pomegranate', 'Dishwasher', 'Crab',
|
||||
'Hoverboard', 'Meat ball', 'Rice Cooker', 'Tuba', 'Calculator', 'Papaya', 'Antelope', 'Parrot', 'Seal',
|
||||
'Butterfly', 'Dumbbell', 'Donkey', 'Lion', 'Urinal', 'Dolphin', 'Electric Drill', 'Hair Dryer', 'Egg tart',
|
||||
'Jellyfish', 'Treadmill', 'Lighter', 'Grapefruit', 'Game board', 'Mop', 'Radish', 'Baozi', 'Target', 'French',
|
||||
'Spring Rolls', 'Monkey', 'Rabbit', 'Pencil Case', 'Yak', 'Red Cabbage', 'Binoculars', 'Asparagus', 'Barbell',
|
||||
'Scallop', 'Noddles', 'Comb', 'Dumpling', 'Oyster', 'Table Tennis paddle', 'Cosmetics Brush/Eyeliner Pencil',
|
||||
'Chainsaw', 'Eraser', 'Lobster', 'Durian', 'Okra', 'Lipstick', 'Cosmetics Mirror', 'Curling', 'Table Tennis' ]
|
||||
|
||||
|
||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
||||
download: |
|
||||
from pycocotools.coco import COCO
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils.general import download, Path
|
||||
|
||||
# Make Directories
|
||||
dir = Path('../datasets/objects365') # dataset directory
|
||||
for p in 'images', 'labels':
|
||||
(dir / p).mkdir(parents=True, exist_ok=True)
|
||||
for q in 'train', 'val':
|
||||
(dir / p / q).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Download
|
||||
url = "https://dorc.ks3-cn-beijing.ksyun.com/data-set/2020Objects365%E6%95%B0%E6%8D%AE%E9%9B%86/train/"
|
||||
download([url + 'zhiyuan_objv2_train.tar.gz'], dir=dir, delete=False) # annotations json
|
||||
download([url + f for f in [f'patch{i}.tar.gz' for i in range(51)]], dir=dir / 'images' / 'train',
|
||||
curl=True, delete=False, threads=8)
|
||||
|
||||
# Move
|
||||
train = dir / 'images' / 'train'
|
||||
for f in tqdm(train.rglob('*.jpg'), desc=f'Moving images'):
|
||||
f.rename(train / f.name) # move to /images/train
|
||||
|
||||
# Labels
|
||||
coco = COCO(dir / 'zhiyuan_objv2_train.json')
|
||||
names = [x["name"] for x in coco.loadCats(coco.getCatIds())]
|
||||
for cid, cat in enumerate(names):
|
||||
catIds = coco.getCatIds(catNms=[cat])
|
||||
imgIds = coco.getImgIds(catIds=catIds)
|
||||
for im in tqdm(coco.loadImgs(imgIds), desc=f'Class {cid + 1}/{len(names)} {cat}'):
|
||||
width, height = im["width"], im["height"]
|
||||
path = Path(im["file_name"]) # image filename
|
||||
try:
|
||||
with open(dir / 'labels' / 'train' / path.with_suffix('.txt').name, 'a') as file:
|
||||
annIds = coco.getAnnIds(imgIds=im["id"], catIds=catIds, iscrowd=None)
|
||||
for a in coco.loadAnns(annIds):
|
||||
x, y, w, h = a['bbox'] # bounding box in xywh (xy top-left corner)
|
||||
x, y = x + w / 2, y + h / 2 # xy to center
|
||||
file.write(f"{cid} {x / width:.5f} {y / height:.5f} {w / width:.5f} {h / height:.5f}\n")
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
@ -2,10 +2,10 @@
|
||||
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
|
||||
# Download command: bash data/scripts/get_argoverse_hd.sh
|
||||
# Train command: python train.py --data argoverse_hd.yaml
|
||||
# Default dataset location is next to /yolov5:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /argoverse
|
||||
# /yolov5
|
||||
# /yolov3
|
||||
|
||||
# Download/unzip images
|
||||
d='../argoverse/' # unzip directory
|
||||
@ -25,7 +25,7 @@ import json
|
||||
from pathlib import Path
|
||||
|
||||
annotation_files = ["train.json", "val.json"]
|
||||
print("Converting annotations to YOLOv5 format...")
|
||||
print("Converting annotations to YOLOv3 format...")
|
||||
|
||||
for val in annotation_files:
|
||||
a = json.load(open(val, "rb"))
|
||||
@ -36,7 +36,7 @@ for val in annotation_files:
|
||||
img_name = a['images'][img_id]['name']
|
||||
img_label_name = img_name[:-3] + "txt"
|
||||
|
||||
obj_class = annot['category_id']
|
||||
cls = annot['category_id'] # instance class id
|
||||
x_center, y_center, width, height = annot['bbox']
|
||||
x_center = (x_center + width / 2) / 1920. # offset and scale
|
||||
y_center = (y_center + height / 2) / 1200. # offset and scale
|
||||
@ -46,11 +46,10 @@ for val in annotation_files:
|
||||
img_dir = "./labels/" + a['seq_dirs'][a['images'][annot['image_id']]['sid']]
|
||||
|
||||
Path(img_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if img_dir + "/" + img_label_name not in label_dict:
|
||||
label_dict[img_dir + "/" + img_label_name] = []
|
||||
|
||||
label_dict[img_dir + "/" + img_label_name].append(f"{obj_class} {x_center} {y_center} {width} {height}\n")
|
||||
label_dict[img_dir + "/" + img_label_name].append(f"{cls} {x_center} {y_center} {width} {height}\n")
|
||||
|
||||
for filename in label_dict:
|
||||
with open(filename, "w") as file:
|
||||
|
||||
@ -2,7 +2,7 @@
|
||||
# COCO 2017 dataset http://cocodataset.org
|
||||
# Download command: bash data/scripts/get_coco.sh
|
||||
# Train command: python train.py --data coco.yaml
|
||||
# Default dataset location is next to /yolov3:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /coco
|
||||
# /yolov3
|
||||
|
||||
17
data/scripts/get_coco128.sh
Normal file
17
data/scripts/get_coco128.sh
Normal file
@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128
|
||||
# Download command: bash data/scripts/get_coco128.sh
|
||||
# Train command: python train.py --data coco128.yaml
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /coco128
|
||||
# /yolov3
|
||||
|
||||
# Download/unzip images and labels
|
||||
d='../' # unzip directory
|
||||
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
||||
f='coco128.zip' # or 'coco2017labels-segments.zip', 68 MB
|
||||
echo 'Downloading' $url$f ' ...'
|
||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
|
||||
|
||||
wait # finish background tasks
|
||||
@ -2,10 +2,10 @@
|
||||
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
|
||||
# Download command: bash data/scripts/get_voc.sh
|
||||
# Train command: python train.py --data voc.yaml
|
||||
# Default dataset location is next to /yolov5:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /VOC
|
||||
# /yolov5
|
||||
# /yolov3
|
||||
|
||||
start=$(date +%s)
|
||||
mkdir -p ../tmp
|
||||
@ -29,34 +29,27 @@ echo "Completed in" $runtime "seconds"
|
||||
|
||||
echo "Splitting dataset..."
|
||||
python3 - "$@" <<END
|
||||
import xml.etree.ElementTree as ET
|
||||
import pickle
|
||||
import os
|
||||
from os import listdir, getcwd
|
||||
from os.path import join
|
||||
import xml.etree.ElementTree as ET
|
||||
from os import getcwd
|
||||
|
||||
sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
|
||||
sets = [('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
|
||||
|
||||
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
|
||||
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
|
||||
"horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
|
||||
|
||||
|
||||
def convert(size, box):
|
||||
dw = 1./(size[0])
|
||||
dh = 1./(size[1])
|
||||
x = (box[0] + box[1])/2.0 - 1
|
||||
y = (box[2] + box[3])/2.0 - 1
|
||||
w = box[1] - box[0]
|
||||
h = box[3] - box[2]
|
||||
x = x*dw
|
||||
w = w*dw
|
||||
y = y*dh
|
||||
h = h*dh
|
||||
return (x,y,w,h)
|
||||
def convert_box(size, box):
|
||||
dw = 1. / (size[0])
|
||||
dh = 1. / (size[1])
|
||||
x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
|
||||
return x * dw, y * dh, w * dw, h * dh
|
||||
|
||||
|
||||
def convert_annotation(year, image_id):
|
||||
in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
|
||||
out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
|
||||
tree=ET.parse(in_file)
|
||||
in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml' % (year, image_id))
|
||||
out_file = open('VOCdevkit/VOC%s/labels/%s.txt' % (year, image_id), 'w')
|
||||
tree = ET.parse(in_file)
|
||||
root = tree.getroot()
|
||||
size = root.find('size')
|
||||
w = int(size.find('width').text)
|
||||
@ -65,74 +58,58 @@ def convert_annotation(year, image_id):
|
||||
for obj in root.iter('object'):
|
||||
difficult = obj.find('difficult').text
|
||||
cls = obj.find('name').text
|
||||
if cls not in classes or int(difficult)==1:
|
||||
if cls not in classes or int(difficult) == 1:
|
||||
continue
|
||||
cls_id = classes.index(cls)
|
||||
xmlbox = obj.find('bndbox')
|
||||
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
|
||||
bb = convert((w,h), b)
|
||||
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
|
||||
float(xmlbox.find('ymax').text))
|
||||
bb = convert_box((w, h), b)
|
||||
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
|
||||
|
||||
wd = getcwd()
|
||||
|
||||
cwd = getcwd()
|
||||
for year, image_set in sets:
|
||||
if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
|
||||
os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
|
||||
image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
|
||||
list_file = open('%s_%s.txt'%(year, image_set), 'w')
|
||||
if not os.path.exists('VOCdevkit/VOC%s/labels/' % year):
|
||||
os.makedirs('VOCdevkit/VOC%s/labels/' % year)
|
||||
image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt' % (year, image_set)).read().strip().split()
|
||||
list_file = open('%s_%s.txt' % (year, image_set), 'w')
|
||||
for image_id in image_ids:
|
||||
list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
|
||||
list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n' % (cwd, year, image_id))
|
||||
convert_annotation(year, image_id)
|
||||
list_file.close()
|
||||
|
||||
END
|
||||
|
||||
cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt >train.txt
|
||||
cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt
|
||||
|
||||
mkdir ../VOC ../VOC/images ../VOC/images/train ../VOC/images/val
|
||||
mkdir ../VOC/labels ../VOC/labels/train ../VOC/labels/val
|
||||
|
||||
python3 - "$@" <<END
|
||||
|
||||
import shutil
|
||||
import os
|
||||
os.system('mkdir ../VOC/')
|
||||
os.system('mkdir ../VOC/images')
|
||||
os.system('mkdir ../VOC/images/train')
|
||||
os.system('mkdir ../VOC/images/val')
|
||||
|
||||
os.system('mkdir ../VOC/labels')
|
||||
os.system('mkdir ../VOC/labels/train')
|
||||
os.system('mkdir ../VOC/labels/val')
|
||||
|
||||
import os
|
||||
print(os.path.exists('../tmp/train.txt'))
|
||||
f = open('../tmp/train.txt', 'r')
|
||||
lines = f.readlines()
|
||||
|
||||
for line in lines:
|
||||
with open('../tmp/train.txt', 'r') as f:
|
||||
for line in f.readlines():
|
||||
line = "/".join(line.split('/')[-5:]).strip()
|
||||
if (os.path.exists("../" + line)):
|
||||
os.system("cp ../"+ line + " ../VOC/images/train")
|
||||
|
||||
line = line.replace('JPEGImages', 'labels')
|
||||
line = line.replace('jpg', 'txt')
|
||||
if (os.path.exists("../" + line)):
|
||||
os.system("cp ../"+ line + " ../VOC/labels/train")
|
||||
if os.path.exists("../" + line):
|
||||
os.system("cp ../" + line + " ../VOC/images/train")
|
||||
|
||||
line = line.replace('JPEGImages', 'labels').replace('jpg', 'txt')
|
||||
if os.path.exists("../" + line):
|
||||
os.system("cp ../" + line + " ../VOC/labels/train")
|
||||
|
||||
print(os.path.exists('../tmp/2007_test.txt'))
|
||||
f = open('../tmp/2007_test.txt', 'r')
|
||||
lines = f.readlines()
|
||||
|
||||
for line in lines:
|
||||
with open('../tmp/2007_test.txt', 'r') as f:
|
||||
for line in f.readlines():
|
||||
line = "/".join(line.split('/')[-5:]).strip()
|
||||
if (os.path.exists("../" + line)):
|
||||
os.system("cp ../"+ line + " ../VOC/images/val")
|
||||
|
||||
line = line.replace('JPEGImages', 'labels')
|
||||
line = line.replace('jpg', 'txt')
|
||||
if (os.path.exists("../" + line)):
|
||||
os.system("cp ../"+ line + " ../VOC/labels/val")
|
||||
if os.path.exists("../" + line):
|
||||
os.system("cp ../" + line + " ../VOC/images/val")
|
||||
|
||||
line = line.replace('JPEGImages', 'labels').replace('jpg', 'txt')
|
||||
if os.path.exists("../" + line):
|
||||
os.system("cp ../" + line + " ../VOC/labels/val")
|
||||
END
|
||||
|
||||
rm -rf ../tmp # remove temporary directory
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
|
||||
# Train command: python train.py --data voc.yaml
|
||||
# Default dataset location is next to /yolov3:
|
||||
# Default dataset location is next to YOLOv3:
|
||||
# /parent_folder
|
||||
# /VOC
|
||||
# /yolov3
|
||||
|
||||
44
detect.py
44
detect.py
@ -5,24 +5,24 @@ from pathlib import Path
|
||||
import cv2
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
from numpy import random
|
||||
|
||||
from models.experimental import attempt_load
|
||||
from utils.datasets import LoadStreams, LoadImages
|
||||
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
|
||||
scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
|
||||
from utils.plots import plot_one_box
|
||||
scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
|
||||
from utils.plots import colors, plot_one_box
|
||||
from utils.torch_utils import select_device, load_classifier, time_synchronized
|
||||
|
||||
|
||||
def detect(save_img=False):
|
||||
@torch.no_grad()
|
||||
def detect(opt):
|
||||
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
||||
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
|
||||
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
|
||||
('rtsp://', 'rtmp://', 'http://', 'https://'))
|
||||
|
||||
# Directories
|
||||
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
|
||||
save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
|
||||
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
||||
|
||||
# Initialize
|
||||
@ -34,6 +34,7 @@ def detect(save_img=False):
|
||||
model = attempt_load(weights, map_location=device) # load FP32 model
|
||||
stride = int(model.stride.max()) # model stride
|
||||
imgsz = check_img_size(imgsz, s=stride) # check img_size
|
||||
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
||||
if half:
|
||||
model.half() # to FP16
|
||||
|
||||
@ -52,10 +53,6 @@ def detect(save_img=False):
|
||||
else:
|
||||
dataset = LoadImages(source, img_size=imgsz, stride=stride)
|
||||
|
||||
# Get names and colors
|
||||
names = model.module.names if hasattr(model, 'module') else model.names
|
||||
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
|
||||
|
||||
# Run inference
|
||||
if device.type != 'cpu':
|
||||
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
|
||||
@ -72,7 +69,8 @@ def detect(save_img=False):
|
||||
pred = model(img, augment=opt.augment)[0]
|
||||
|
||||
# Apply NMS
|
||||
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
|
||||
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms,
|
||||
max_det=opt.max_det)
|
||||
t2 = time_synchronized()
|
||||
|
||||
# Apply Classifier
|
||||
@ -82,15 +80,16 @@ def detect(save_img=False):
|
||||
# Process detections
|
||||
for i, det in enumerate(pred): # detections per image
|
||||
if webcam: # batch_size >= 1
|
||||
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
|
||||
p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
|
||||
else:
|
||||
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
|
||||
p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
|
||||
|
||||
p = Path(p) # to Path
|
||||
save_path = str(save_dir / p.name) # img.jpg
|
||||
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
|
||||
s += '%gx%g ' % img.shape[2:] # print string
|
||||
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
|
||||
imc = im0.copy() if opt.save_crop else im0 # for opt.save_crop
|
||||
if len(det):
|
||||
# Rescale boxes from img_size to im0 size
|
||||
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
||||
@ -108,9 +107,12 @@ def detect(save_img=False):
|
||||
with open(txt_path + '.txt', 'a') as f:
|
||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
||||
|
||||
if save_img or view_img: # Add bbox to image
|
||||
label = f'{names[int(cls)]} {conf:.2f}'
|
||||
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
|
||||
if save_img or opt.save_crop or view_img: # Add bbox to image
|
||||
c = int(cls) # integer class
|
||||
label = None if opt.hide_labels else (names[c] if opt.hide_conf else f'{names[c]} {conf:.2f}')
|
||||
plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=opt.line_thickness)
|
||||
if opt.save_crop:
|
||||
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
|
||||
|
||||
# Print time (inference + NMS)
|
||||
print(f'{s}Done. ({t2 - t1:.3f}s)')
|
||||
@ -153,10 +155,12 @@ if __name__ == '__main__':
|
||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
||||
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
|
||||
parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--view-img', action='store_true', help='display results')
|
||||
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
||||
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
||||
parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
|
||||
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
|
||||
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
|
||||
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
|
||||
@ -165,14 +169,16 @@ if __name__ == '__main__':
|
||||
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
|
||||
parser.add_argument('--name', default='exp', help='save results to project/name')
|
||||
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
||||
parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
|
||||
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
|
||||
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
|
||||
opt = parser.parse_args()
|
||||
print(opt)
|
||||
check_requirements(exclude=('pycocotools', 'thop'))
|
||||
check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
|
||||
|
||||
with torch.no_grad():
|
||||
if opt.update: # update all models (to fix SourceChangeWarning)
|
||||
for opt.weights in ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']:
|
||||
detect()
|
||||
detect(opt=opt)
|
||||
strip_optimizer(opt.weights)
|
||||
else:
|
||||
detect()
|
||||
detect(opt=opt)
|
||||
|
||||
97
hubconf.py
97
hubconf.py
@ -2,24 +2,13 @@
|
||||
|
||||
Usage:
|
||||
import torch
|
||||
model = torch.hub.load('ultralytics/yolov3', 'yolov3tiny')
|
||||
model = torch.hub.load('ultralytics/yolov3', 'yolov3_tiny')
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from models.yolo import Model
|
||||
from utils.general import check_requirements, set_logging
|
||||
from utils.google_utils import attempt_download
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
dependencies = ['torch', 'yaml']
|
||||
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
|
||||
set_logging()
|
||||
|
||||
|
||||
def create(name, pretrained, channels, classes, autoshape):
|
||||
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||
"""Creates a specified YOLOv3 model
|
||||
|
||||
Arguments:
|
||||
@ -27,17 +16,32 @@ def create(name, pretrained, channels, classes, autoshape):
|
||||
pretrained (bool): load pretrained weights into the model
|
||||
channels (int): number of input channels
|
||||
classes (int): number of model classes
|
||||
autoshape (bool): apply YOLOv3 .autoshape() wrapper to model
|
||||
verbose (bool): print all information to screen
|
||||
device (str, torch.device, None): device to use for model parameters
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
YOLOv3 pytorch model
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
from models.yolo import Model, attempt_load
|
||||
from utils.general import check_requirements, set_logging
|
||||
from utils.google_utils import attempt_download
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop'))
|
||||
set_logging(verbose=verbose)
|
||||
|
||||
fname = Path(name).with_suffix('.pt') # checkpoint filename
|
||||
try:
|
||||
if pretrained and channels == 3 and classes == 80:
|
||||
model = attempt_load(fname, map_location=torch.device('cpu')) # download/load FP32 model
|
||||
else:
|
||||
cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
|
||||
model = Model(cfg, channels, classes)
|
||||
model = Model(cfg, channels, classes) # create model
|
||||
if pretrained:
|
||||
fname = f'{name}.pt' # checkpoint filename
|
||||
attempt_download(fname) # download if not found locally
|
||||
ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
|
||||
ckpt = torch.load(attempt_download(fname), map_location=torch.device('cpu')) # load
|
||||
msd = model.state_dict() # model state_dict
|
||||
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
|
||||
csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter
|
||||
@ -46,66 +50,47 @@ def create(name, pretrained, channels, classes, autoshape):
|
||||
model.names = ckpt['model'].names # set class names attribute
|
||||
if autoshape:
|
||||
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
|
||||
device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
|
||||
device = select_device('0' if torch.cuda.is_available() else 'cpu') if device is None else torch.device(device)
|
||||
return model.to(device)
|
||||
|
||||
except Exception as e:
|
||||
help_url = 'https://github.com/ultralytics/yolov5/issues/36'
|
||||
s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url
|
||||
s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url
|
||||
raise Exception(s) from e
|
||||
|
||||
|
||||
def custom(path_or_model='path/to/model.pt', autoshape=True):
|
||||
"""YOLOv3-custom model https://github.com/ultralytics/yolov3
|
||||
|
||||
Arguments (3 options):
|
||||
path_or_model (str): 'path/to/model.pt'
|
||||
path_or_model (dict): torch.load('path/to/model.pt')
|
||||
path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
|
||||
if isinstance(model, dict):
|
||||
model = model['ema' if model.get('ema') else 'model'] # load model
|
||||
|
||||
hub_model = Model(model.yaml).to(next(model.parameters()).device) # create
|
||||
hub_model.load_state_dict(model.float().state_dict()) # load state_dict
|
||||
hub_model.names = model.names # class names
|
||||
if autoshape:
|
||||
hub_model = hub_model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
|
||||
device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
|
||||
return hub_model.to(device)
|
||||
def custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None):
|
||||
# YOLOv3 custom or local model
|
||||
return _create(path, autoshape=autoshape, verbose=verbose, device=device)
|
||||
|
||||
|
||||
def yolov3(pretrained=True, channels=3, classes=80, autoshape=True):
|
||||
def yolov3(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||
# YOLOv3 model https://github.com/ultralytics/yolov3
|
||||
return create('yolov3', pretrained, channels, classes, autoshape)
|
||||
return _create('yolov3', pretrained, channels, classes, autoshape, verbose, device)
|
||||
|
||||
|
||||
def yolov3_spp(pretrained=True, channels=3, classes=80, autoshape=True):
|
||||
def yolov3_spp(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||
# YOLOv3-SPP model https://github.com/ultralytics/yolov3
|
||||
return create('yolov3-spp', pretrained, channels, classes, autoshape)
|
||||
return _create('yolov3-spp', pretrained, channels, classes, autoshape, verbose, device)
|
||||
|
||||
|
||||
def yolov3_tiny(pretrained=True, channels=3, classes=80, autoshape=True):
|
||||
def yolov3_tiny(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||
# YOLOv3-tiny model https://github.com/ultralytics/yolov3
|
||||
return create('yolov3-tiny', pretrained, channels, classes, autoshape)
|
||||
return _create('yolov3-tiny', pretrained, channels, classes, autoshape, verbose, device)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = create(name='yolov3', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example
|
||||
# model = custom(path_or_model='path/to/model.pt') # custom example
|
||||
model = _create(name='yolov3', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
|
||||
# model = custom(path='path/to/model.pt') # custom
|
||||
|
||||
# Verify inference
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
imgs = [Image.open('data/images/bus.jpg'), # PIL
|
||||
'data/images/zidane.jpg', # filename
|
||||
'https://github.com/ultralytics/yolov3/raw/master/data/images/bus.jpg', # URI
|
||||
np.zeros((640, 480, 3))] # numpy
|
||||
imgs = ['data/images/zidane.jpg', # filename
|
||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg', # URI
|
||||
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
|
||||
Image.open('data/images/bus.jpg'), # PIL
|
||||
np.zeros((320, 640, 3))] # numpy
|
||||
|
||||
results = model(imgs) # batched inference
|
||||
results.print()
|
||||
|
||||
@ -13,8 +13,8 @@ from PIL import Image
|
||||
from torch.cuda import amp
|
||||
|
||||
from utils.datasets import letterbox
|
||||
from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh
|
||||
from utils.plots import color_list, plot_one_box
|
||||
from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh, save_one_box
|
||||
from utils.plots import colors, plot_one_box
|
||||
from utils.torch_utils import time_synchronized
|
||||
|
||||
|
||||
@ -215,32 +215,34 @@ class NMS(nn.Module):
|
||||
conf = 0.25 # confidence threshold
|
||||
iou = 0.45 # IoU threshold
|
||||
classes = None # (optional list) filter by class
|
||||
max_det = 1000 # maximum number of detections per image
|
||||
|
||||
def __init__(self):
|
||||
super(NMS, self).__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
|
||||
return non_max_suppression(x[0], self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det)
|
||||
|
||||
|
||||
class autoShape(nn.Module):
|
||||
class AutoShape(nn.Module):
|
||||
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
|
||||
conf = 0.25 # NMS confidence threshold
|
||||
iou = 0.45 # NMS IoU threshold
|
||||
classes = None # (optional list) filter by class
|
||||
max_det = 1000 # maximum number of detections per image
|
||||
|
||||
def __init__(self, model):
|
||||
super(autoShape, self).__init__()
|
||||
super(AutoShape, self).__init__()
|
||||
self.model = model.eval()
|
||||
|
||||
def autoshape(self):
|
||||
print('autoShape already enabled, skipping... ') # model already converted to model.autoshape()
|
||||
print('AutoShape already enabled, skipping... ') # model already converted to model.autoshape()
|
||||
return self
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, imgs, size=640, augment=False, profile=False):
|
||||
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
||||
# filename: imgs = 'data/samples/zidane.jpg'
|
||||
# filename: imgs = 'data/images/zidane.jpg'
|
||||
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
|
||||
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
|
||||
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
|
||||
@ -271,7 +273,7 @@ class autoShape(nn.Module):
|
||||
shape0.append(s) # image shape
|
||||
g = (size / max(s)) # gain
|
||||
shape1.append([y * g for y in s])
|
||||
imgs[i] = im # update
|
||||
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
|
||||
shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
|
||||
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
|
||||
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
|
||||
@ -285,7 +287,7 @@ class autoShape(nn.Module):
|
||||
t.append(time_synchronized())
|
||||
|
||||
# Post-process
|
||||
y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
|
||||
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det) # NMS
|
||||
for i in range(n):
|
||||
scale_coords(shape1, y[i][:, :4], shape0[i])
|
||||
|
||||
@ -311,29 +313,32 @@ class Detections:
|
||||
self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
|
||||
self.s = shape # inference BCHW shape
|
||||
|
||||
def display(self, pprint=False, show=False, save=False, render=False, save_dir=''):
|
||||
colors = color_list()
|
||||
for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
|
||||
str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
|
||||
def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
|
||||
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
|
||||
str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
|
||||
if pred is not None:
|
||||
for c in pred[:, -1].unique():
|
||||
n = (pred[:, -1] == c).sum() # detections per class
|
||||
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
||||
if show or save or render:
|
||||
if show or save or render or crop:
|
||||
for *box, conf, cls in pred: # xyxy, confidence, class
|
||||
label = f'{self.names[int(cls)]} {conf:.2f}'
|
||||
plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
|
||||
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
|
||||
if crop:
|
||||
save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i])
|
||||
else: # all others
|
||||
plot_one_box(box, im, label=label, color=colors(cls))
|
||||
|
||||
im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
|
||||
if pprint:
|
||||
print(str.rstrip(', '))
|
||||
if show:
|
||||
img.show(self.files[i]) # show
|
||||
im.show(self.files[i]) # show
|
||||
if save:
|
||||
f = self.files[i]
|
||||
img.save(Path(save_dir) / f) # save
|
||||
im.save(save_dir / f) # save
|
||||
print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')
|
||||
if render:
|
||||
self.imgs[i] = np.asarray(img)
|
||||
self.imgs[i] = np.asarray(im)
|
||||
|
||||
def print(self):
|
||||
self.display(pprint=True) # print results
|
||||
@ -343,10 +348,14 @@ class Detections:
|
||||
self.display(show=True) # show results
|
||||
|
||||
def save(self, save_dir='runs/hub/exp'):
|
||||
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp') # increment save_dir
|
||||
Path(save_dir).mkdir(parents=True, exist_ok=True)
|
||||
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir
|
||||
self.display(save=True, save_dir=save_dir) # save results
|
||||
|
||||
def crop(self, save_dir='runs/hub/exp'):
|
||||
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir
|
||||
self.display(crop=True, save_dir=save_dir) # crop results
|
||||
print(f'Saved results to {save_dir}\n')
|
||||
|
||||
def render(self):
|
||||
self.display(render=True) # render results
|
||||
return self.imgs
|
||||
|
||||
@ -110,25 +110,27 @@ class Ensemble(nn.ModuleList):
|
||||
return y, None # inference, train output
|
||||
|
||||
|
||||
def attempt_load(weights, map_location=None):
|
||||
def attempt_load(weights, map_location=None, inplace=True):
|
||||
from models.yolo import Detect, Model
|
||||
|
||||
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
|
||||
model = Ensemble()
|
||||
for w in weights if isinstance(weights, list) else [weights]:
|
||||
attempt_download(w)
|
||||
ckpt = torch.load(w, map_location=map_location) # load
|
||||
ckpt = torch.load(attempt_download(w), map_location=map_location) # load
|
||||
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model
|
||||
|
||||
# Compatibility updates
|
||||
for m in model.modules():
|
||||
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
|
||||
m.inplace = True # pytorch 1.7.0 compatibility
|
||||
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model]:
|
||||
m.inplace = inplace # pytorch 1.7.0 compatibility
|
||||
elif type(m) is Conv:
|
||||
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
||||
|
||||
if len(model) == 1:
|
||||
return model[-1] # return model
|
||||
else:
|
||||
print('Ensemble created with %s\n' % weights)
|
||||
for k in ['names', 'stride']:
|
||||
print(f'Ensemble created with {weights}\n')
|
||||
for k in ['names']:
|
||||
setattr(model, k, getattr(model[-1], k))
|
||||
model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride
|
||||
return model # return ensemble
|
||||
|
||||
105
models/export.py
105
models/export.py
@ -1,34 +1,43 @@
|
||||
"""Exports a YOLOv3 *.pt model to ONNX and TorchScript formats
|
||||
"""Exports a YOLOv3 *.pt model to TorchScript, ONNX, CoreML formats
|
||||
|
||||
Usage:
|
||||
$ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov3.pt --img 640 --batch 1
|
||||
$ python path/to/models/export.py --weights yolov3.pt --img 640 --batch 1
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
||||
sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.mobile_optimizer import optimize_for_mobile
|
||||
|
||||
import models
|
||||
from models.experimental import attempt_load
|
||||
from utils.activations import Hardswish, SiLU
|
||||
from utils.general import set_logging, check_img_size
|
||||
from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', type=str, default='./yolov3.pt', help='weights path') # from yolov3/models/
|
||||
parser.add_argument('--weights', type=str, default='./yolov3.pt', help='weights path')
|
||||
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
|
||||
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
||||
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
|
||||
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
|
||||
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
|
||||
parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
|
||||
parser.add_argument('--inplace', action='store_true', help='set YOLOv3 Detect() inplace=True')
|
||||
parser.add_argument('--train', action='store_true', help='model.train() mode')
|
||||
parser.add_argument('--optimize', action='store_true', help='optimize TorchScript for mobile') # TorchScript-only
|
||||
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only
|
||||
parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
|
||||
parser.add_argument('--opset-version', type=int, default=12, help='ONNX opset version') # ONNX-only
|
||||
opt = parser.parse_args()
|
||||
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
|
||||
opt.include = [x.lower() for x in opt.include]
|
||||
print(opt)
|
||||
set_logging()
|
||||
t = time.time()
|
||||
@ -41,11 +50,16 @@ if __name__ == '__main__':
|
||||
# Checks
|
||||
gs = int(max(model.stride)) # grid size (max stride)
|
||||
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
|
||||
assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
|
||||
|
||||
# Input
|
||||
img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection
|
||||
|
||||
# Update model
|
||||
if opt.half:
|
||||
img, model = img.half(), model.half() # to FP16
|
||||
if opt.train:
|
||||
model.train() # training mode (no grid construction in Detect layer)
|
||||
for k, m in model.named_modules():
|
||||
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
||||
if isinstance(m, models.common.Conv): # assign export-friendly activations
|
||||
@ -53,52 +67,79 @@ if __name__ == '__main__':
|
||||
m.act = Hardswish()
|
||||
elif isinstance(m.act, nn.SiLU):
|
||||
m.act = SiLU()
|
||||
# elif isinstance(m, models.yolo.Detect):
|
||||
elif isinstance(m, models.yolo.Detect):
|
||||
m.inplace = opt.inplace
|
||||
m.onnx_dynamic = opt.dynamic
|
||||
# m.forward = m.forward_export # assign forward (optional)
|
||||
model.model[-1].export = not opt.grid # set Detect() layer grid export
|
||||
y = model(img) # dry run
|
||||
|
||||
# TorchScript export
|
||||
for _ in range(2):
|
||||
y = model(img) # dry runs
|
||||
print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
|
||||
|
||||
# TorchScript export -----------------------------------------------------------------------------------------------
|
||||
if 'torchscript' in opt.include or 'coreml' in opt.include:
|
||||
prefix = colorstr('TorchScript:')
|
||||
try:
|
||||
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
|
||||
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
||||
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
||||
ts = torch.jit.trace(model, img, strict=False)
|
||||
ts.save(f)
|
||||
print('TorchScript export success, saved as %s' % f)
|
||||
(optimize_for_mobile(ts) if opt.optimize else ts).save(f)
|
||||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||
except Exception as e:
|
||||
print('TorchScript export failure: %s' % e)
|
||||
print(f'{prefix} export failure: {e}')
|
||||
|
||||
# ONNX export
|
||||
# ONNX export ------------------------------------------------------------------------------------------------------
|
||||
if 'onnx' in opt.include:
|
||||
prefix = colorstr('ONNX:')
|
||||
try:
|
||||
import onnx
|
||||
|
||||
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
|
||||
print(f'{prefix} starting export with onnx {onnx.__version__}...')
|
||||
f = opt.weights.replace('.pt', '.onnx') # filename
|
||||
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
|
||||
output_names=['classes', 'boxes'] if y is None else ['output'],
|
||||
torch.onnx.export(model, img, f, verbose=False, opset_version=opt.opset_version, input_names=['images'],
|
||||
training=torch.onnx.TrainingMode.TRAINING if opt.train else torch.onnx.TrainingMode.EVAL,
|
||||
do_constant_folding=not opt.train,
|
||||
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
|
||||
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
|
||||
|
||||
# Checks
|
||||
onnx_model = onnx.load(f) # load onnx model
|
||||
onnx.checker.check_model(onnx_model) # check onnx model
|
||||
# print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
|
||||
print('ONNX export success, saved as %s' % f)
|
||||
except Exception as e:
|
||||
print('ONNX export failure: %s' % e)
|
||||
model_onnx = onnx.load(f) # load onnx model
|
||||
onnx.checker.check_model(model_onnx) # check onnx model
|
||||
# print(onnx.helper.printable_graph(model_onnx.graph)) # print
|
||||
|
||||
# CoreML export
|
||||
# Simplify
|
||||
if opt.simplify:
|
||||
try:
|
||||
check_requirements(['onnx-simplifier'])
|
||||
import onnxsim
|
||||
|
||||
print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
|
||||
model_onnx, check = onnxsim.simplify(
|
||||
model_onnx,
|
||||
dynamic_input_shape=opt.dynamic,
|
||||
input_shapes={'images': list(img.shape)} if opt.dynamic else None)
|
||||
assert check, 'assert check failed'
|
||||
onnx.save(model_onnx, f)
|
||||
except Exception as e:
|
||||
print(f'{prefix} simplifier failure: {e}')
|
||||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||
except Exception as e:
|
||||
print(f'{prefix} export failure: {e}')
|
||||
|
||||
# CoreML export ----------------------------------------------------------------------------------------------------
|
||||
if 'coreml' in opt.include:
|
||||
prefix = colorstr('CoreML:')
|
||||
try:
|
||||
import coremltools as ct
|
||||
|
||||
print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
|
||||
# convert model from torchscript and apply pixel scaling as per detect.py
|
||||
model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
||||
print(f'{prefix} starting export with coremltools {ct.__version__}...')
|
||||
assert opt.train, 'CoreML exports should be placed in model.train() mode with `python export.py --train`'
|
||||
model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
||||
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
||||
model.save(f)
|
||||
print('CoreML export success, saved as %s' % f)
|
||||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||
except Exception as e:
|
||||
print('CoreML export failure: %s' % e)
|
||||
print(f'{prefix} export failure: {e}')
|
||||
|
||||
# Finish
|
||||
print('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t))
|
||||
print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.')
|
||||
|
||||
@ -1,11 +1,16 @@
|
||||
# YOLOv3 YOLO-specific modules
|
||||
"""YOLOv3-specific modules
|
||||
|
||||
Usage:
|
||||
$ python path/to/models/yolo.py --cfg yolov3.yaml
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
||||
sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from models.common import *
|
||||
@ -23,9 +28,9 @@ except ImportError:
|
||||
|
||||
class Detect(nn.Module):
|
||||
stride = None # strides computed during build
|
||||
export = False # onnx export
|
||||
onnx_dynamic = False # ONNX export parameter
|
||||
|
||||
def __init__(self, nc=80, anchors=(), ch=()): # detection layer
|
||||
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
|
||||
super(Detect, self).__init__()
|
||||
self.nc = nc # number of classes
|
||||
self.no = nc + 5 # number of outputs per anchor
|
||||
@ -36,23 +41,28 @@ class Detect(nn.Module):
|
||||
self.register_buffer('anchors', a) # shape(nl,na,2)
|
||||
self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2)
|
||||
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
|
||||
self.inplace = inplace # use in-place ops (e.g. slice assignment)
|
||||
|
||||
def forward(self, x):
|
||||
# x = x.copy() # for profiling
|
||||
z = [] # inference output
|
||||
self.training |= self.export
|
||||
for i in range(self.nl):
|
||||
x[i] = self.m[i](x[i]) # conv
|
||||
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
|
||||
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
||||
|
||||
if not self.training: # inference
|
||||
if self.grid[i].shape[2:4] != x[i].shape[2:4]:
|
||||
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
|
||||
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
|
||||
|
||||
y = x[i].sigmoid()
|
||||
if self.inplace:
|
||||
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
|
||||
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
|
||||
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
|
||||
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
|
||||
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
|
||||
y = torch.cat((xy, wh, y[..., 4:]), -1)
|
||||
z.append(y.view(bs, -1, self.no))
|
||||
|
||||
return x if self.training else (torch.cat(z, 1), x)
|
||||
@ -72,7 +82,7 @@ class Model(nn.Module):
|
||||
import yaml # for torch hub
|
||||
self.yaml_file = Path(cfg).name
|
||||
with open(cfg) as f:
|
||||
self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict
|
||||
self.yaml = yaml.safe_load(f) # model dict
|
||||
|
||||
# Define model
|
||||
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
||||
@ -84,18 +94,20 @@ class Model(nn.Module):
|
||||
self.yaml['anchors'] = round(anchors) # override yaml value
|
||||
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
|
||||
self.names = [str(i) for i in range(self.yaml['nc'])] # default names
|
||||
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
||||
self.inplace = self.yaml.get('inplace', True)
|
||||
# logger.info([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
||||
|
||||
# Build strides, anchors
|
||||
m = self.model[-1] # Detect()
|
||||
if isinstance(m, Detect):
|
||||
s = 256 # 2x min stride
|
||||
m.inplace = self.inplace
|
||||
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
|
||||
m.anchors /= m.stride.view(-1, 1, 1)
|
||||
check_anchor_order(m)
|
||||
self.stride = m.stride
|
||||
self._initialize_biases() # only run once
|
||||
# print('Strides: %s' % m.stride.tolist())
|
||||
# logger.info('Strides: %s' % m.stride.tolist())
|
||||
|
||||
# Init weights, biases
|
||||
initialize_weights(self)
|
||||
@ -104,6 +116,11 @@ class Model(nn.Module):
|
||||
|
||||
def forward(self, x, augment=False, profile=False):
|
||||
if augment:
|
||||
return self.forward_augment(x) # augmented inference, None
|
||||
else:
|
||||
return self.forward_once(x, profile) # single-scale inference, train
|
||||
|
||||
def forward_augment(self, x):
|
||||
img_size = x.shape[-2:] # height, width
|
||||
s = [1, 0.83, 0.67] # scales
|
||||
f = [None, 3, None] # flips (2-ud, 3-lr)
|
||||
@ -112,15 +129,9 @@ class Model(nn.Module):
|
||||
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
|
||||
yi = self.forward_once(xi)[0] # forward
|
||||
# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
|
||||
yi[..., :4] /= si # de-scale
|
||||
if fi == 2:
|
||||
yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud
|
||||
elif fi == 3:
|
||||
yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr
|
||||
yi = self._descale_pred(yi, fi, si, img_size)
|
||||
y.append(yi)
|
||||
return torch.cat(y, 1), None # augmented inference, train
|
||||
else:
|
||||
return self.forward_once(x, profile) # single-scale inference, train
|
||||
|
||||
def forward_once(self, x, profile=False):
|
||||
y, dt = [], [] # outputs
|
||||
@ -134,15 +145,34 @@ class Model(nn.Module):
|
||||
for _ in range(10):
|
||||
_ = m(x)
|
||||
dt.append((time_synchronized() - t) * 100)
|
||||
print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
|
||||
if m == self.model[0]:
|
||||
logger.info(f"{'time (ms)':>10s} {'GFLOPS':>10s} {'params':>10s} {'module'}")
|
||||
logger.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}')
|
||||
|
||||
x = m(x) # run
|
||||
y.append(x if m.i in self.save else None) # save output
|
||||
|
||||
if profile:
|
||||
print('%.1fms total' % sum(dt))
|
||||
logger.info('%.1fms total' % sum(dt))
|
||||
return x
|
||||
|
||||
def _descale_pred(self, p, flips, scale, img_size):
|
||||
# de-scale predictions following augmented inference (inverse operation)
|
||||
if self.inplace:
|
||||
p[..., :4] /= scale # de-scale
|
||||
if flips == 2:
|
||||
p[..., 1] = img_size[0] - p[..., 1] # de-flip ud
|
||||
elif flips == 3:
|
||||
p[..., 0] = img_size[1] - p[..., 0] # de-flip lr
|
||||
else:
|
||||
x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale # de-scale
|
||||
if flips == 2:
|
||||
y = img_size[0] - y # de-flip ud
|
||||
elif flips == 3:
|
||||
x = img_size[1] - x # de-flip lr
|
||||
p = torch.cat((x, y, wh, p[..., 4:]), -1)
|
||||
return p
|
||||
|
||||
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
|
||||
# https://arxiv.org/abs/1708.02002 section 3.3
|
||||
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
|
||||
@ -157,15 +187,16 @@ class Model(nn.Module):
|
||||
m = self.model[-1] # Detect() module
|
||||
for mi in m.m: # from
|
||||
b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
|
||||
print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
|
||||
logger.info(
|
||||
('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
|
||||
|
||||
# def _print_weights(self):
|
||||
# for m in self.model.modules():
|
||||
# if type(m) is Bottleneck:
|
||||
# print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
|
||||
# logger.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
|
||||
|
||||
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
||||
print('Fusing layers... ')
|
||||
logger.info('Fusing layers... ')
|
||||
for m in self.model.modules():
|
||||
if type(m) is Conv and hasattr(m, 'bn'):
|
||||
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
||||
@ -177,20 +208,20 @@ class Model(nn.Module):
|
||||
def nms(self, mode=True): # add or remove NMS module
|
||||
present = type(self.model[-1]) is NMS # last layer is NMS
|
||||
if mode and not present:
|
||||
print('Adding NMS... ')
|
||||
logger.info('Adding NMS... ')
|
||||
m = NMS() # module
|
||||
m.f = -1 # from
|
||||
m.i = self.model[-1].i + 1 # index
|
||||
self.model.add_module(name='%s' % m.i, module=m) # add
|
||||
self.eval()
|
||||
elif not mode and present:
|
||||
print('Removing NMS... ')
|
||||
logger.info('Removing NMS... ')
|
||||
self.model = self.model[:-1] # remove
|
||||
return self
|
||||
|
||||
def autoshape(self): # add autoShape module
|
||||
print('Adding autoShape... ')
|
||||
m = autoShape(self) # wrap model
|
||||
def autoshape(self): # add AutoShape module
|
||||
logger.info('Adding AutoShape... ')
|
||||
m = AutoShape(self) # wrap model
|
||||
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
|
||||
return m
|
||||
|
||||
@ -266,12 +297,12 @@ if __name__ == '__main__':
|
||||
model.train()
|
||||
|
||||
# Profile
|
||||
# img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
|
||||
# img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 320, 320).to(device)
|
||||
# y = model(img, profile=True)
|
||||
|
||||
# Tensorboard
|
||||
# Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
|
||||
# from torch.utils.tensorboard import SummaryWriter
|
||||
# tb_writer = SummaryWriter()
|
||||
# print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/")
|
||||
# tb_writer.add_graph(model.model, img) # add model to tensorboard
|
||||
# tb_writer = SummaryWriter('.')
|
||||
# logger.info("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/")
|
||||
# tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph
|
||||
# tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard
|
||||
|
||||
@ -21,9 +21,10 @@ pandas
|
||||
|
||||
# export --------------------------------------
|
||||
# coremltools>=4.1
|
||||
# onnx>=1.8.1
|
||||
# onnx>=1.9.0
|
||||
# scikit-learn==0.19.2 # for coreml quantization
|
||||
|
||||
# extras --------------------------------------
|
||||
thop # FLOPS computation
|
||||
# Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
|
||||
pycocotools>=2.0 # COCO mAP
|
||||
thop # FLOPS computation
|
||||
|
||||
24
test.py
24
test.py
@ -18,6 +18,7 @@ from utils.plots import plot_images, output_to_target, plot_study_txt
|
||||
from utils.torch_utils import select_device, time_synchronized
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test(data,
|
||||
weights=None,
|
||||
batch_size=32,
|
||||
@ -38,7 +39,8 @@ def test(data,
|
||||
wandb_logger=None,
|
||||
compute_loss=None,
|
||||
half_precision=True,
|
||||
is_coco=False):
|
||||
is_coco=False,
|
||||
opt=None):
|
||||
# Initialize/load model and set device
|
||||
training = model is not None
|
||||
if training: # called by train.py
|
||||
@ -49,7 +51,7 @@ def test(data,
|
||||
device = select_device(opt.device, batch_size=batch_size)
|
||||
|
||||
# Directories
|
||||
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
|
||||
save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
|
||||
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
||||
|
||||
# Load model
|
||||
@ -71,7 +73,7 @@ def test(data,
|
||||
if isinstance(data, str):
|
||||
is_coco = data.endswith('coco.yaml')
|
||||
with open(data) as f:
|
||||
data = yaml.load(f, Loader=yaml.SafeLoader)
|
||||
data = yaml.safe_load(f)
|
||||
check_dataset(data) # check
|
||||
nc = 1 if single_cls else int(data['nc']) # number of classes
|
||||
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
|
||||
@ -104,7 +106,6 @@ def test(data,
|
||||
targets = targets.to(device)
|
||||
nb, _, height, width = img.shape # batch size, channels, height, width
|
||||
|
||||
with torch.no_grad():
|
||||
# Run model
|
||||
t = time_synchronized()
|
||||
out, train_out = model(img, augment=augment) # inference and training outputs
|
||||
@ -118,7 +119,7 @@ def test(data,
|
||||
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
||||
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
||||
t = time_synchronized()
|
||||
out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=True)
|
||||
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
|
||||
t1 += time_synchronized() - t
|
||||
|
||||
# Statistics per image
|
||||
@ -135,6 +136,8 @@ def test(data,
|
||||
continue
|
||||
|
||||
# Predictions
|
||||
if single_cls:
|
||||
pred[:, 5] = 0
|
||||
predn = pred.clone()
|
||||
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
|
||||
|
||||
@ -185,8 +188,8 @@ def test(data,
|
||||
|
||||
# Per target class
|
||||
for cls in torch.unique(tcls_tensor):
|
||||
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices
|
||||
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices
|
||||
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # target indices
|
||||
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # prediction indices
|
||||
|
||||
# Search for detections
|
||||
if pi.shape[0]:
|
||||
@ -307,7 +310,7 @@ if __name__ == '__main__':
|
||||
opt.save_json |= opt.data.endswith('coco.yaml')
|
||||
opt.data = check_file(opt.data) # check file
|
||||
print(opt)
|
||||
check_requirements()
|
||||
check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
|
||||
|
||||
if opt.task in ('train', 'val', 'test'): # run normally
|
||||
test(opt.data,
|
||||
@ -323,11 +326,12 @@ if __name__ == '__main__':
|
||||
save_txt=opt.save_txt | opt.save_hybrid,
|
||||
save_hybrid=opt.save_hybrid,
|
||||
save_conf=opt.save_conf,
|
||||
opt=opt
|
||||
)
|
||||
|
||||
elif opt.task == 'speed': # speed benchmarks
|
||||
for w in opt.weights:
|
||||
test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False)
|
||||
test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False, opt=opt)
|
||||
|
||||
elif opt.task == 'study': # run over a range of settings and save/plot
|
||||
# python test.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt yolov3-tiny.pt
|
||||
@ -338,7 +342,7 @@ if __name__ == '__main__':
|
||||
for i in x: # img-size
|
||||
print(f'\nRunning {f} point {i}...')
|
||||
r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
|
||||
plots=False)
|
||||
plots=False, opt=opt)
|
||||
y.append(r + t) # results and times
|
||||
np.savetxt(f, y, fmt='%10.4g') # save
|
||||
os.system('zip -r study.zip study_*.txt')
|
||||
|
||||
54
train.py
54
train.py
@ -32,7 +32,7 @@ from utils.general import labels_to_class_weights, increment_path, labels_to_ima
|
||||
from utils.google_utils import attempt_download
|
||||
from utils.loss import ComputeLoss
|
||||
from utils.plots import plot_images, plot_labels, plot_results, plot_evolution
|
||||
from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel
|
||||
from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, de_parallel
|
||||
from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -52,24 +52,23 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
|
||||
# Save run settings
|
||||
with open(save_dir / 'hyp.yaml', 'w') as f:
|
||||
yaml.dump(hyp, f, sort_keys=False)
|
||||
yaml.safe_dump(hyp, f, sort_keys=False)
|
||||
with open(save_dir / 'opt.yaml', 'w') as f:
|
||||
yaml.dump(vars(opt), f, sort_keys=False)
|
||||
yaml.safe_dump(vars(opt), f, sort_keys=False)
|
||||
|
||||
# Configure
|
||||
plots = not opt.evolve # create plots
|
||||
cuda = device.type != 'cpu'
|
||||
init_seeds(2 + rank)
|
||||
with open(opt.data) as f:
|
||||
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
|
||||
is_coco = opt.data.endswith('coco.yaml')
|
||||
data_dict = yaml.safe_load(f) # data dict
|
||||
|
||||
# Logging- Doing this before checking the dataset. Might update data_dict
|
||||
loggers = {'wandb': None} # loggers dict
|
||||
if rank in [-1, 0]:
|
||||
opt.hyp = hyp # add hyperparameters
|
||||
run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
|
||||
wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
|
||||
wandb_logger = WandbLogger(opt, save_dir.stem, run_id, data_dict)
|
||||
loggers['wandb'] = wandb_logger.wandb
|
||||
data_dict = wandb_logger.data_dict
|
||||
if wandb_logger.wandb:
|
||||
@ -78,12 +77,13 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
|
||||
names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
|
||||
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
|
||||
is_coco = opt.data.endswith('coco.yaml') and nc == 80 # COCO dataset
|
||||
|
||||
# Model
|
||||
pretrained = weights.endswith('.pt')
|
||||
if pretrained:
|
||||
with torch_distributed_zero_first(rank):
|
||||
attempt_download(weights) # download if not found locally
|
||||
weights = attempt_download(weights) # download if not found locally
|
||||
ckpt = torch.load(weights, map_location=device) # load checkpoint
|
||||
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
|
||||
exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys
|
||||
@ -330,9 +330,9 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
if plots and ni < 3:
|
||||
f = save_dir / f'train_batch{ni}.jpg' # filename
|
||||
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
|
||||
# if tb_writer:
|
||||
if tb_writer:
|
||||
tb_writer.add_graph(torch.jit.trace(de_parallel(model), imgs, strict=False), []) # model graph
|
||||
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
|
||||
# tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
|
||||
elif plots and ni == 10 and wandb_logger.wandb:
|
||||
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
|
||||
save_dir.glob('train*.jpg') if x.exists()]})
|
||||
@ -358,6 +358,7 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
single_cls=opt.single_cls,
|
||||
dataloader=testloader,
|
||||
save_dir=save_dir,
|
||||
save_json=is_coco and final_epoch,
|
||||
verbose=nc < 50 and final_epoch,
|
||||
plots=plots and final_epoch,
|
||||
wandb_logger=wandb_logger,
|
||||
@ -367,8 +368,6 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
# Write
|
||||
with open(results_file, 'a') as f:
|
||||
f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss
|
||||
if len(opt.name) and opt.bucket:
|
||||
os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
|
||||
|
||||
# Log
|
||||
tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
|
||||
@ -392,7 +391,7 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
ckpt = {'epoch': epoch,
|
||||
'best_fitness': best_fitness,
|
||||
'training_results': results_file.read_text(),
|
||||
'model': deepcopy(model.module if is_parallel(model) else model).half(),
|
||||
'model': deepcopy(de_parallel(model)).half(),
|
||||
'ema': deepcopy(ema.ema).half(),
|
||||
'updates': ema.updates,
|
||||
'optimizer': optimizer.state_dict(),
|
||||
@ -411,17 +410,17 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
# end epoch ----------------------------------------------------------------------------------------------------
|
||||
# end training
|
||||
if rank in [-1, 0]:
|
||||
# Plots
|
||||
logger.info(f'{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.\n')
|
||||
if plots:
|
||||
plot_results(save_dir=save_dir) # save as results.png
|
||||
if wandb_logger.wandb:
|
||||
files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
|
||||
wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files
|
||||
if (save_dir / f).exists()]})
|
||||
# Test best.pt
|
||||
logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
|
||||
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
|
||||
for m in (last, best) if best.exists() else (last): # speed, mAP tests
|
||||
|
||||
if not opt.evolve:
|
||||
if is_coco: # COCO dataset
|
||||
for m in [last, best] if best.exists() else [last]: # speed, mAP tests
|
||||
results, _, _ = test.test(opt.data,
|
||||
batch_size=batch_size * 2,
|
||||
imgsz=imgsz_test,
|
||||
@ -436,16 +435,13 @@ def train(hyp, opt, device, tb_writer=None):
|
||||
is_coco=is_coco)
|
||||
|
||||
# Strip optimizers
|
||||
final = best if best.exists() else last # final model
|
||||
for f in last, best:
|
||||
if f.exists():
|
||||
strip_optimizer(f) # strip optimizers
|
||||
if opt.bucket:
|
||||
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
|
||||
if wandb_logger.wandb and not opt.evolve: # Log the stripped model
|
||||
wandb_logger.wandb.log_artifact(str(final), type='model',
|
||||
if wandb_logger.wandb: # Log the stripped model
|
||||
wandb_logger.wandb.log_artifact(str(best if best.exists() else last), type='model',
|
||||
name='run_' + wandb_logger.wandb_run.id + '_model',
|
||||
aliases=['last', 'best', 'stripped'])
|
||||
aliases=['latest', 'best', 'stripped'])
|
||||
wandb_logger.finish_run()
|
||||
else:
|
||||
dist.destroy_process_group()
|
||||
@ -497,7 +493,7 @@ if __name__ == '__main__':
|
||||
set_logging(opt.global_rank)
|
||||
if opt.global_rank in [-1, 0]:
|
||||
check_git_status()
|
||||
check_requirements()
|
||||
check_requirements(exclude=('pycocotools', 'thop'))
|
||||
|
||||
# Resume
|
||||
wandb_run = check_wandb_resume(opt)
|
||||
@ -506,8 +502,9 @@ if __name__ == '__main__':
|
||||
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
|
||||
apriori = opt.global_rank, opt.local_rank
|
||||
with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
|
||||
opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader)) # replace
|
||||
opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori # reinstate
|
||||
opt = argparse.Namespace(**yaml.safe_load(f)) # replace
|
||||
opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = \
|
||||
'', ckpt, True, opt.total_batch_size, *apriori # reinstate
|
||||
logger.info('Resuming training from %s' % ckpt)
|
||||
else:
|
||||
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
|
||||
@ -515,7 +512,7 @@ if __name__ == '__main__':
|
||||
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
|
||||
opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test)
|
||||
opt.name = 'evolve' if opt.evolve else opt.name
|
||||
opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve) # increment run
|
||||
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve))
|
||||
|
||||
# DDP mode
|
||||
opt.total_batch_size = opt.batch_size
|
||||
@ -526,11 +523,12 @@ if __name__ == '__main__':
|
||||
device = torch.device('cuda', opt.local_rank)
|
||||
dist.init_process_group(backend='nccl', init_method='env://') # distributed backend
|
||||
assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
|
||||
assert not opt.image_weights, '--image-weights argument is not compatible with DDP training'
|
||||
opt.batch_size = opt.total_batch_size // opt.world_size
|
||||
|
||||
# Hyperparameters
|
||||
with open(opt.hyp) as f:
|
||||
hyp = yaml.load(f, Loader=yaml.SafeLoader) # load hyps
|
||||
hyp = yaml.safe_load(f) # load hyps
|
||||
|
||||
# Train
|
||||
logger.info(opt)
|
||||
|
||||
@ -19,23 +19,6 @@ class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
|
||||
return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX
|
||||
|
||||
|
||||
class MemoryEfficientSwish(nn.Module):
|
||||
class F(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x):
|
||||
ctx.save_for_backward(x)
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x = ctx.saved_tensors[0]
|
||||
sx = torch.sigmoid(x)
|
||||
return grad_output * (sx * (1 + x * (1 - sx)))
|
||||
|
||||
def forward(self, x):
|
||||
return self.F.apply(x)
|
||||
|
||||
|
||||
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
||||
class Mish(nn.Module):
|
||||
@staticmethod
|
||||
@ -70,3 +53,46 @@ class FReLU(nn.Module):
|
||||
|
||||
def forward(self, x):
|
||||
return torch.max(x, self.bn(self.conv(x)))
|
||||
|
||||
|
||||
# ACON https://arxiv.org/pdf/2009.04759.pdf ----------------------------------------------------------------------------
|
||||
class AconC(nn.Module):
|
||||
r""" ACON activation (activate or not).
|
||||
AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
|
||||
according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
|
||||
"""
|
||||
|
||||
def __init__(self, c1):
|
||||
super().__init__()
|
||||
self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
|
||||
self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1))
|
||||
self.beta = nn.Parameter(torch.ones(1, c1, 1, 1))
|
||||
|
||||
def forward(self, x):
|
||||
dpx = (self.p1 - self.p2) * x
|
||||
return dpx * torch.sigmoid(self.beta * dpx) + self.p2 * x
|
||||
|
||||
|
||||
class MetaAconC(nn.Module):
|
||||
r""" ACON activation (activate or not).
|
||||
MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network
|
||||
according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
|
||||
"""
|
||||
|
||||
def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r
|
||||
super().__init__()
|
||||
c2 = max(r, c1 // r)
|
||||
self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1))
|
||||
self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1))
|
||||
self.fc1 = nn.Conv2d(c1, c2, k, s, bias=True)
|
||||
self.fc2 = nn.Conv2d(c2, c1, k, s, bias=True)
|
||||
# self.bn1 = nn.BatchNorm2d(c2)
|
||||
# self.bn2 = nn.BatchNorm2d(c1)
|
||||
|
||||
def forward(self, x):
|
||||
y = x.mean(dim=2, keepdims=True).mean(dim=3, keepdims=True)
|
||||
# batch-size 1 bug/instabilities https://github.com/ultralytics/yolov5/issues/2891
|
||||
# beta = torch.sigmoid(self.bn2(self.fc2(self.bn1(self.fc1(y))))) # bug/unstable
|
||||
beta = torch.sigmoid(self.fc2(self.fc1(y))) # bug patch BN layers removed
|
||||
dpx = (self.p1 - self.p2) * x
|
||||
return dpx * torch.sigmoid(beta * dpx) + self.p2 * x
|
||||
|
||||
@ -3,7 +3,6 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
from scipy.cluster.vq import kmeans
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils.general import colorstr
|
||||
@ -76,6 +75,8 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
||||
Usage:
|
||||
from utils.autoanchor import *; _ = kmean_anchors()
|
||||
"""
|
||||
from scipy.cluster.vq import kmeans
|
||||
|
||||
thr = 1. / thr
|
||||
prefix = colorstr('autoanchor: ')
|
||||
|
||||
@ -102,7 +103,7 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
||||
|
||||
if isinstance(path, str): # *.yaml file
|
||||
with open(path) as f:
|
||||
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # model dict
|
||||
data_dict = yaml.safe_load(f) # model dict
|
||||
from utils.datasets import LoadImagesAndLabels
|
||||
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
|
||||
else:
|
||||
|
||||
@ -19,7 +19,7 @@ for last in path.rglob('*/**/last.pt'):
|
||||
|
||||
# Load opt.yaml
|
||||
with open(last.parent.parent / 'opt.yaml') as f:
|
||||
opt = yaml.load(f, Loader=yaml.SafeLoader)
|
||||
opt = yaml.safe_load(f)
|
||||
|
||||
# Get device count
|
||||
d = opt['device'].split(',') # devices
|
||||
|
||||
@ -7,7 +7,7 @@
|
||||
cd home/ubuntu
|
||||
if [ ! -d yolov5 ]; then
|
||||
echo "Running first-time script." # install dependencies, download COCO, pull Docker
|
||||
git clone https://github.com/ultralytics/yolov5 && sudo chmod -R 777 yolov5
|
||||
git clone https://github.com/ultralytics/yolov5 -b master && sudo chmod -R 777 yolov5
|
||||
cd yolov5
|
||||
bash data/scripts/get_coco.sh && echo "Data done." &
|
||||
sudo docker pull ultralytics/yolov5:latest && echo "Docker done." &
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
# Dataset utils and dataloaders
|
||||
|
||||
import glob
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
@ -36,9 +37,12 @@ for orientation in ExifTags.TAGS.keys():
|
||||
break
|
||||
|
||||
|
||||
def get_hash(files):
|
||||
# Returns a single hash value of a list of files
|
||||
return sum(os.path.getsize(f) for f in files if os.path.isfile(f))
|
||||
def get_hash(paths):
|
||||
# Returns a single hash value of a list of paths (files or dirs)
|
||||
size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes
|
||||
h = hashlib.md5(str(size).encode()) # hash sizes
|
||||
h.update(''.join(paths).encode()) # hash paths
|
||||
return h.hexdigest() # return hash
|
||||
|
||||
|
||||
def exif_size(img):
|
||||
@ -172,12 +176,12 @@ class LoadImages: # for inference
|
||||
ret_val, img0 = self.cap.read()
|
||||
|
||||
self.frame += 1
|
||||
print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='')
|
||||
print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ', end='')
|
||||
|
||||
else:
|
||||
# Read image
|
||||
self.count += 1
|
||||
img0 = cv2.imread(path) # BGR
|
||||
img0 = cv2.imread(path, -1) # BGR (-1 is IMREAD_UNCHANGED)
|
||||
assert img0 is not None, 'Image Not Found ' + path
|
||||
print(f'image {self.count}/{self.nf} {path}: ', end='')
|
||||
|
||||
@ -193,7 +197,7 @@ class LoadImages: # for inference
|
||||
def new_video(self, path):
|
||||
self.frame = 0
|
||||
self.cap = cv2.VideoCapture(path)
|
||||
self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
def __len__(self):
|
||||
return self.nf # number of files
|
||||
@ -270,26 +274,27 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||
sources = [sources]
|
||||
|
||||
n = len(sources)
|
||||
self.imgs = [None] * n
|
||||
self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
|
||||
self.sources = [clean_str(x) for x in sources] # clean source names for later
|
||||
for i, s in enumerate(sources):
|
||||
# Start the thread to read frames from the video stream
|
||||
for i, s in enumerate(sources): # index, source
|
||||
# Start thread to read frames from video stream
|
||||
print(f'{i + 1}/{n}: {s}... ', end='')
|
||||
url = eval(s) if s.isnumeric() else s
|
||||
if 'youtube.com/' in url or 'youtu.be/' in url: # if source is YouTube video
|
||||
if 'youtube.com/' in s or 'youtu.be/' in s: # if source is YouTube video
|
||||
check_requirements(('pafy', 'youtube_dl'))
|
||||
import pafy
|
||||
url = pafy.new(url).getbest(preftype="mp4").url
|
||||
cap = cv2.VideoCapture(url)
|
||||
s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL
|
||||
s = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam
|
||||
cap = cv2.VideoCapture(s)
|
||||
assert cap.isOpened(), f'Failed to open {s}'
|
||||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
self.fps = cap.get(cv2.CAP_PROP_FPS) % 100
|
||||
self.fps[i] = max(cap.get(cv2.CAP_PROP_FPS) % 100, 0) or 30.0 # 30 FPS fallback
|
||||
self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback
|
||||
|
||||
_, self.imgs[i] = cap.read() # guarantee first frame
|
||||
thread = Thread(target=self.update, args=([i, cap]), daemon=True)
|
||||
print(f' success ({w}x{h} at {self.fps:.2f} FPS).')
|
||||
thread.start()
|
||||
self.threads[i] = Thread(target=self.update, args=([i, cap]), daemon=True)
|
||||
print(f" success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")
|
||||
self.threads[i].start()
|
||||
print('') # newline
|
||||
|
||||
# check for common shapes
|
||||
@ -298,18 +303,17 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||
if not self.rect:
|
||||
print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
|
||||
|
||||
def update(self, index, cap):
|
||||
# Read next stream frame in a daemon thread
|
||||
n = 0
|
||||
while cap.isOpened():
|
||||
def update(self, i, cap):
|
||||
# Read stream `i` frames in daemon thread
|
||||
n, f = 0, self.frames[i]
|
||||
while cap.isOpened() and n < f:
|
||||
n += 1
|
||||
# _, self.imgs[index] = cap.read()
|
||||
cap.grab()
|
||||
if n == 4: # read every 4th frame
|
||||
if n % 4: # read every 4th frame
|
||||
success, im = cap.retrieve()
|
||||
self.imgs[index] = im if success else self.imgs[index] * 0
|
||||
n = 0
|
||||
time.sleep(1 / self.fps) # wait time
|
||||
self.imgs[i] = im if success else self.imgs[i] * 0
|
||||
time.sleep(1 / self.fps[i]) # wait time
|
||||
|
||||
def __iter__(self):
|
||||
self.count = -1
|
||||
@ -317,12 +321,12 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||
|
||||
def __next__(self):
|
||||
self.count += 1
|
||||
img0 = self.imgs.copy()
|
||||
if cv2.waitKey(1) == ord('q'): # q to quit
|
||||
if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit
|
||||
cv2.destroyAllWindows()
|
||||
raise StopIteration
|
||||
|
||||
# Letterbox
|
||||
img0 = self.imgs.copy()
|
||||
img = [letterbox(x, self.img_size, auto=self.rect, stride=self.stride)[0] for x in img0]
|
||||
|
||||
# Stack
|
||||
@ -383,7 +387,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') # cached labels
|
||||
if cache_path.is_file():
|
||||
cache, exists = torch.load(cache_path), True # load
|
||||
if cache['hash'] != get_hash(self.label_files + self.img_files) or 'version' not in cache: # changed
|
||||
if cache['hash'] != get_hash(self.label_files + self.img_files): # changed
|
||||
cache, exists = self.cache_labels(cache_path, prefix), False # re-cache
|
||||
else:
|
||||
cache, exists = self.cache_labels(cache_path, prefix), False # cache
|
||||
@ -470,7 +474,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
if os.path.isfile(lb_file):
|
||||
nf += 1 # label found
|
||||
with open(lb_file, 'r') as f:
|
||||
l = [x.split() for x in f.read().strip().splitlines()]
|
||||
l = [x.split() for x in f.read().strip().splitlines() if len(x)]
|
||||
if any([len(x) > 8 for x in l]): # is segment
|
||||
classes = np.array([x[0] for x in l], dtype=np.float32)
|
||||
segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
|
||||
@ -490,20 +494,23 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
x[im_file] = [l, shape, segments]
|
||||
except Exception as e:
|
||||
nc += 1
|
||||
print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}')
|
||||
logging.info(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}')
|
||||
|
||||
pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels... " \
|
||||
f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
|
||||
pbar.close()
|
||||
|
||||
if nf == 0:
|
||||
print(f'{prefix}WARNING: No labels found in {path}. See {help_url}')
|
||||
logging.info(f'{prefix}WARNING: No labels found in {path}. See {help_url}')
|
||||
|
||||
x['hash'] = get_hash(self.label_files + self.img_files)
|
||||
x['results'] = nf, nm, ne, nc, i + 1
|
||||
x['version'] = 0.1 # cache version
|
||||
torch.save(x, path) # save for next time
|
||||
x['version'] = 0.2 # cache version
|
||||
try:
|
||||
torch.save(x, path) # save cache for next time
|
||||
logging.info(f'{prefix}New cache created: {path}')
|
||||
except Exception as e:
|
||||
logging.info(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # path not writeable
|
||||
return x
|
||||
|
||||
def __len__(self):
|
||||
@ -634,10 +641,10 @@ def load_image(self, index):
|
||||
img = cv2.imread(path) # BGR
|
||||
assert img is not None, 'Image Not Found ' + path
|
||||
h0, w0 = img.shape[:2] # orig hw
|
||||
r = self.img_size / max(h0, w0) # resize image to img_size
|
||||
if r != 1: # always resize down, only resize up if training with augmentation
|
||||
interp = cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR
|
||||
img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
|
||||
r = self.img_size / max(h0, w0) # ratio
|
||||
if r != 1: # if sizes are not equal
|
||||
img = cv2.resize(img, (int(w0 * r), int(h0 * r)),
|
||||
interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR)
|
||||
return img, (h0, w0), img.shape[:2] # img, hw_original, hw_resized
|
||||
else:
|
||||
return self.imgs[index], self.img_hw0[index], self.img_hw[index] # img, hw_original, hw_resized
|
||||
|
||||
68
utils/flask_rest_api/README.md
Normal file
68
utils/flask_rest_api/README.md
Normal file
@ -0,0 +1,68 @@
|
||||
# Flask REST API
|
||||
[REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/).
|
||||
|
||||
## Requirements
|
||||
|
||||
[Flask](https://palletsprojects.com/p/flask/) is required. Install with:
|
||||
```shell
|
||||
$ pip install Flask
|
||||
```
|
||||
|
||||
## Run
|
||||
|
||||
After Flask installation run:
|
||||
|
||||
```shell
|
||||
$ python3 restapi.py --port 5000
|
||||
```
|
||||
|
||||
Then use [curl](https://curl.se/) to perform a request:
|
||||
|
||||
```shell
|
||||
$ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s'`
|
||||
```
|
||||
|
||||
The model inference results are returned as a JSON response:
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"class": 0,
|
||||
"confidence": 0.8900438547,
|
||||
"height": 0.9318675399,
|
||||
"name": "person",
|
||||
"width": 0.3264600933,
|
||||
"xcenter": 0.7438579798,
|
||||
"ycenter": 0.5207948685
|
||||
},
|
||||
{
|
||||
"class": 0,
|
||||
"confidence": 0.8440024257,
|
||||
"height": 0.7155083418,
|
||||
"name": "person",
|
||||
"width": 0.6546785235,
|
||||
"xcenter": 0.427829951,
|
||||
"ycenter": 0.6334488392
|
||||
},
|
||||
{
|
||||
"class": 27,
|
||||
"confidence": 0.3771208823,
|
||||
"height": 0.3902671337,
|
||||
"name": "tie",
|
||||
"width": 0.0696444362,
|
||||
"xcenter": 0.3675483763,
|
||||
"ycenter": 0.7991207838
|
||||
},
|
||||
{
|
||||
"class": 27,
|
||||
"confidence": 0.3527112305,
|
||||
"height": 0.1540903747,
|
||||
"name": "tie",
|
||||
"width": 0.0336618312,
|
||||
"xcenter": 0.7814827561,
|
||||
"ycenter": 0.5065554976
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given in `example_request.py`
|
||||
13
utils/flask_rest_api/example_request.py
Normal file
13
utils/flask_rest_api/example_request.py
Normal file
@ -0,0 +1,13 @@
|
||||
"""Perform test request"""
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
|
||||
TEST_IMAGE = "zidane.jpg"
|
||||
|
||||
image_data = open(TEST_IMAGE, "rb").read()
|
||||
|
||||
response = requests.post(DETECTION_URL, files={"image": image_data}).json()
|
||||
|
||||
pprint.pprint(response)
|
||||
37
utils/flask_rest_api/restapi.py
Normal file
37
utils/flask_rest_api/restapi.py
Normal file
@ -0,0 +1,37 @@
|
||||
"""
|
||||
Run a rest API exposing the yolov5s object detection model
|
||||
"""
|
||||
import argparse
|
||||
import io
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from flask import Flask, request
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
DETECTION_URL = "/v1/object-detection/yolov5s"
|
||||
|
||||
|
||||
@app.route(DETECTION_URL, methods=["POST"])
|
||||
def predict():
|
||||
if not request.method == "POST":
|
||||
return
|
||||
|
||||
if request.files.get("image"):
|
||||
image_file = request.files["image"]
|
||||
image_bytes = image_file.read()
|
||||
|
||||
img = Image.open(io.BytesIO(image_bytes))
|
||||
|
||||
results = model(img, size=640) # reduce size=320 for faster inference
|
||||
return results.pandas().xyxy[0].to_json(orient="records")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Flask API exposing YOLOv3 model")
|
||||
parser.add_argument("--port", default=5000, type=int, help="port number")
|
||||
args = parser.parse_args()
|
||||
|
||||
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force_reload to recache
|
||||
app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat
|
||||
147
utils/general.py
147
utils/general.py
@ -9,11 +9,14 @@ import random
|
||||
import re
|
||||
import subprocess
|
||||
import time
|
||||
from itertools import repeat
|
||||
from multiprocessing.pool import ThreadPool
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pkg_resources as pkg
|
||||
import torch
|
||||
import torchvision
|
||||
import yaml
|
||||
@ -30,10 +33,10 @@ cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with Py
|
||||
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
|
||||
|
||||
|
||||
def set_logging(rank=-1):
|
||||
def set_logging(rank=-1, verbose=True):
|
||||
logging.basicConfig(
|
||||
format="%(message)s",
|
||||
level=logging.INFO if rank in [-1, 0] else logging.WARN)
|
||||
level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
|
||||
|
||||
|
||||
def init_seeds(seed=0):
|
||||
@ -49,16 +52,30 @@ def get_latest_run(search_dir='.'):
|
||||
return max(last_list, key=os.path.getctime) if last_list else ''
|
||||
|
||||
|
||||
def isdocker():
|
||||
def is_docker():
|
||||
# Is environment a Docker container
|
||||
return Path('/workspace').exists() # or Path('/.dockerenv').exists()
|
||||
|
||||
|
||||
def is_colab():
|
||||
# Is environment a Google Colab instance
|
||||
try:
|
||||
import google.colab
|
||||
return True
|
||||
except Exception as e:
|
||||
return False
|
||||
|
||||
|
||||
def emojis(str=''):
|
||||
# Return platform-dependent emoji-safe version of string
|
||||
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
||||
|
||||
|
||||
def file_size(file):
|
||||
# Return file size in MB
|
||||
return Path(file).stat().st_size / 1e6
|
||||
|
||||
|
||||
def check_online():
|
||||
# Check internet connectivity
|
||||
import socket
|
||||
@ -74,7 +91,7 @@ def check_git_status():
|
||||
print(colorstr('github: '), end='')
|
||||
try:
|
||||
assert Path('.git').exists(), 'skipping check (not a git repository)'
|
||||
assert not isdocker(), 'skipping check (Docker image)'
|
||||
assert not is_docker(), 'skipping check (Docker image)'
|
||||
assert check_online(), 'skipping check (offline)'
|
||||
|
||||
cmd = 'git fetch && git config --get remote.origin.url'
|
||||
@ -91,10 +108,19 @@ def check_git_status():
|
||||
print(e)
|
||||
|
||||
|
||||
def check_python(minimum='3.7.0', required=True):
|
||||
# Check current python version vs. required python version
|
||||
current = platform.python_version()
|
||||
result = pkg.parse_version(current) >= pkg.parse_version(minimum)
|
||||
if required:
|
||||
assert result, f'Python {minimum} required by YOLOv3, but Python {current} is currently installed'
|
||||
return result
|
||||
|
||||
|
||||
def check_requirements(requirements='requirements.txt', exclude=()):
|
||||
# Check installed dependencies meet requirements (pass *.txt file or list of packages)
|
||||
import pkg_resources as pkg
|
||||
prefix = colorstr('red', 'bold', 'requirements:')
|
||||
check_python() # check python version
|
||||
if isinstance(requirements, (str, Path)): # requirements.txt file
|
||||
file = Path(requirements)
|
||||
if not file.exists():
|
||||
@ -110,8 +136,11 @@ def check_requirements(requirements='requirements.txt', exclude=()):
|
||||
pkg.require(r)
|
||||
except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
|
||||
n += 1
|
||||
print(f"{prefix} {e.req} not found and is required by YOLOv3, attempting auto-update...")
|
||||
print(subprocess.check_output(f"pip install '{e.req}'", shell=True).decode())
|
||||
print(f"{prefix} {r} not found and is required by YOLOv3, attempting auto-update...")
|
||||
try:
|
||||
print(subprocess.check_output(f"pip install '{r}'", shell=True).decode())
|
||||
except Exception as e:
|
||||
print(f'{prefix} {e}')
|
||||
|
||||
if n: # if packages updated
|
||||
source = file.resolve() if 'file' in locals() else requirements
|
||||
@ -131,7 +160,8 @@ def check_img_size(img_size, s=32):
|
||||
def check_imshow():
|
||||
# Check if environment supports image displays
|
||||
try:
|
||||
assert not isdocker(), 'cv2.imshow() is disabled in Docker environments'
|
||||
assert not is_docker(), 'cv2.imshow() is disabled in Docker environments'
|
||||
assert not is_colab(), 'cv2.imshow() is disabled in Google Colab environments'
|
||||
cv2.imshow('test', np.zeros((1, 1, 3)))
|
||||
cv2.waitKey(1)
|
||||
cv2.destroyAllWindows()
|
||||
@ -143,12 +173,19 @@ def check_imshow():
|
||||
|
||||
|
||||
def check_file(file):
|
||||
# Search for file if not found
|
||||
if Path(file).is_file() or file == '':
|
||||
# Search/download file (if necessary) and return path
|
||||
file = str(file) # convert to str()
|
||||
if Path(file).is_file() or file == '': # exists
|
||||
return file
|
||||
else:
|
||||
elif file.startswith(('http://', 'https://')): # download
|
||||
url, file = file, Path(file).name
|
||||
print(f'Downloading {url} to {file}...')
|
||||
torch.hub.download_url_to_file(url, file)
|
||||
assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
|
||||
return file
|
||||
else: # search
|
||||
files = glob.glob('./**/' + file, recursive=True) # find file
|
||||
assert len(files), f'File Not Found: {file}' # assert file was found
|
||||
assert len(files), f'File not found: {file}' # assert file was found
|
||||
assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
|
||||
return files[0] # return file
|
||||
|
||||
@ -161,18 +198,54 @@ def check_dataset(dict):
|
||||
if not all(x.exists() for x in val):
|
||||
print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
|
||||
if s and len(s): # download script
|
||||
print('Downloading %s ...' % s)
|
||||
if s.startswith('http') and s.endswith('.zip'): # URL
|
||||
f = Path(s).name # filename
|
||||
print(f'Downloading {s} ...')
|
||||
torch.hub.download_url_to_file(s, f)
|
||||
r = os.system('unzip -q %s -d ../ && rm %s' % (f, f)) # unzip
|
||||
else: # bash script
|
||||
r = os.system(f'unzip -q {f} -d ../ && rm {f}') # unzip
|
||||
elif s.startswith('bash '): # bash script
|
||||
print(f'Running {s} ...')
|
||||
r = os.system(s)
|
||||
print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure')) # analyze return value
|
||||
else: # python script
|
||||
r = exec(s) # return None
|
||||
print('Dataset autodownload %s\n' % ('success' if r in (0, None) else 'failure')) # print result
|
||||
else:
|
||||
raise Exception('Dataset not found.')
|
||||
|
||||
|
||||
def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
|
||||
# Multi-threaded file download and unzip function
|
||||
def download_one(url, dir):
|
||||
# Download 1 file
|
||||
f = dir / Path(url).name # filename
|
||||
if not f.exists():
|
||||
print(f'Downloading {url} to {f}...')
|
||||
if curl:
|
||||
os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -") # curl download, retry and resume on fail
|
||||
else:
|
||||
torch.hub.download_url_to_file(url, f, progress=True) # torch download
|
||||
if unzip and f.suffix in ('.zip', '.gz'):
|
||||
print(f'Unzipping {f}...')
|
||||
if f.suffix == '.zip':
|
||||
s = f'unzip -qo {f} -d {dir} && rm {f}' # unzip -quiet -overwrite
|
||||
elif f.suffix == '.gz':
|
||||
s = f'tar xfz {f} --directory {f.parent}' # unzip
|
||||
if delete: # delete zip file after unzip
|
||||
s += f' && rm {f}'
|
||||
os.system(s)
|
||||
|
||||
dir = Path(dir)
|
||||
dir.mkdir(parents=True, exist_ok=True) # make directory
|
||||
if threads > 1:
|
||||
pool = ThreadPool(threads)
|
||||
pool.imap(lambda x: download_one(*x), zip(url, repeat(dir))) # multi-threaded
|
||||
pool.close()
|
||||
pool.join()
|
||||
else:
|
||||
for u in tuple(url) if isinstance(url, str) else url:
|
||||
download_one(u, dir)
|
||||
|
||||
|
||||
def make_divisible(x, divisor):
|
||||
# Returns x evenly divisible by divisor
|
||||
return math.ceil(x / divisor) * divisor
|
||||
@ -419,7 +492,7 @@ def wh_iou(wh1, wh2):
|
||||
|
||||
|
||||
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
||||
labels=()):
|
||||
labels=(), max_det=300):
|
||||
"""Runs Non-Maximum Suppression (NMS) on inference results
|
||||
|
||||
Returns:
|
||||
@ -429,9 +502,12 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
|
||||
nc = prediction.shape[2] - 5 # number of classes
|
||||
xc = prediction[..., 4] > conf_thres # candidates
|
||||
|
||||
# Checks
|
||||
assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
|
||||
assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
|
||||
|
||||
# Settings
|
||||
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
||||
max_det = 300 # maximum number of detections per image
|
||||
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
||||
time_limit = 10.0 # seconds to quit after
|
||||
redundant = True # require redundant detections
|
||||
@ -550,14 +626,14 @@ def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):
|
||||
results = tuple(x[0, :7])
|
||||
c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
|
||||
f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n')
|
||||
yaml.dump(hyp, f, sort_keys=False)
|
||||
yaml.safe_dump(hyp, f, sort_keys=False)
|
||||
|
||||
if bucket:
|
||||
os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket)) # upload
|
||||
|
||||
|
||||
def apply_classifier(x, model, img, im0):
|
||||
# applies a second stage classifier to yolo outputs
|
||||
# Apply a second stage classifier to yolo outputs
|
||||
im0 = [im0] if isinstance(im0, np.ndarray) else im0
|
||||
for i, d in enumerate(x): # per image
|
||||
if d is not None and len(d):
|
||||
@ -591,14 +667,33 @@ def apply_classifier(x, model, img, im0):
|
||||
return x
|
||||
|
||||
|
||||
def increment_path(path, exist_ok=True, sep=''):
|
||||
# Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.
|
||||
def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
|
||||
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
|
||||
xyxy = torch.tensor(xyxy).view(-1, 4)
|
||||
b = xyxy2xywh(xyxy) # boxes
|
||||
if square:
|
||||
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
||||
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
|
||||
xyxy = xywh2xyxy(b).long()
|
||||
clip_coords(xyxy, im.shape)
|
||||
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
|
||||
if save:
|
||||
cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop)
|
||||
return crop
|
||||
|
||||
|
||||
def increment_path(path, exist_ok=False, sep='', mkdir=False):
|
||||
# Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
|
||||
path = Path(path) # os-agnostic
|
||||
if (path.exists() and exist_ok) or (not path.exists()):
|
||||
return str(path)
|
||||
else:
|
||||
if path.exists() and not exist_ok:
|
||||
suffix = path.suffix
|
||||
path = path.with_suffix('')
|
||||
dirs = glob.glob(f"{path}{sep}*") # similar paths
|
||||
matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
|
||||
i = [int(m.groups()[0]) for m in matches if m] # indices
|
||||
n = max(i) + 1 if i else 2 # increment number
|
||||
return f"{path}{sep}{n}" # update path
|
||||
path = Path(f"{path}{sep}{n}{suffix}") # update path
|
||||
dir = path if path.suffix == '' else path.parent # directory
|
||||
if not dir.exists() and mkdir:
|
||||
dir.mkdir(parents=True, exist_ok=True) # make directory
|
||||
return path
|
||||
|
||||
@ -16,40 +16,57 @@ def gsutil_getsize(url=''):
|
||||
return eval(s.split(' ')[0]) if len(s) else 0 # bytes
|
||||
|
||||
|
||||
def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
|
||||
# Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
|
||||
file = Path(file)
|
||||
try: # GitHub
|
||||
print(f'Downloading {url} to {file}...')
|
||||
torch.hub.download_url_to_file(url, str(file))
|
||||
assert file.exists() and file.stat().st_size > min_bytes # check
|
||||
except Exception as e: # GCP
|
||||
file.unlink(missing_ok=True) # remove partial downloads
|
||||
print(f'Download error: {e}\nRe-attempting {url2 or url} to {file}...')
|
||||
os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail
|
||||
finally:
|
||||
if not file.exists() or file.stat().st_size < min_bytes: # check
|
||||
file.unlink(missing_ok=True) # remove partial downloads
|
||||
print(f'ERROR: Download failure: {error_msg or url}')
|
||||
print('')
|
||||
|
||||
|
||||
def attempt_download(file, repo='ultralytics/yolov3'):
|
||||
# Attempt file download if does not exist
|
||||
file = Path(str(file).strip().replace("'", '').lower())
|
||||
file = Path(str(file).strip().replace("'", ''))
|
||||
|
||||
if not file.exists():
|
||||
# URL specified
|
||||
name = file.name
|
||||
if str(file).startswith(('http:/', 'https:/')): # download
|
||||
url = str(file).replace(':/', '://') # Pathlib turns :// -> :/
|
||||
safe_download(file=name, url=url, min_bytes=1E5)
|
||||
return name
|
||||
|
||||
# GitHub assets
|
||||
file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required)
|
||||
try:
|
||||
response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api
|
||||
assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
|
||||
tag = response['tag_name'] # i.e. 'v1.0'
|
||||
except: # fallback plan
|
||||
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
|
||||
tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
|
||||
try:
|
||||
tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
|
||||
except:
|
||||
tag = 'v9.5.0' # current release
|
||||
|
||||
name = file.name
|
||||
if name in assets:
|
||||
msg = f'{file} missing, try downloading from https://github.com/{repo}/releases/'
|
||||
redundant = False # second download option
|
||||
try: # GitHub
|
||||
url = f'https://github.com/{repo}/releases/download/{tag}/{name}'
|
||||
print(f'Downloading {url} to {file}...')
|
||||
torch.hub.download_url_to_file(url, file)
|
||||
assert file.exists() and file.stat().st_size > 1E6 # check
|
||||
except Exception as e: # GCP
|
||||
print(f'Download error: {e}')
|
||||
assert redundant, 'No secondary mirror'
|
||||
url = f'https://storage.googleapis.com/{repo}/ckpt/{name}'
|
||||
print(f'Downloading {url} to {file}...')
|
||||
os.system(f'curl -L {url} -o {file}') # torch.hub.download_url_to_file(url, weights)
|
||||
finally:
|
||||
if not file.exists() or file.stat().st_size < 1E6: # check
|
||||
file.unlink(missing_ok=True) # remove partial downloads
|
||||
print(f'ERROR: Download failure: {msg}')
|
||||
print('')
|
||||
return
|
||||
safe_download(file,
|
||||
url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
|
||||
# url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
|
||||
min_bytes=1E5,
|
||||
error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
|
||||
|
||||
return str(file)
|
||||
|
||||
|
||||
def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
|
||||
|
||||
@ -145,7 +145,7 @@ class ConfusionMatrix:
|
||||
for i, gc in enumerate(gt_classes):
|
||||
j = m0 == i
|
||||
if n and sum(j) == 1:
|
||||
self.matrix[gc, detection_classes[m1[j]]] += 1 # correct
|
||||
self.matrix[detection_classes[m1[j]], gc] += 1 # correct
|
||||
else:
|
||||
self.matrix[self.nc, gc] += 1 # background FP
|
||||
|
||||
|
||||
@ -16,7 +16,6 @@ import seaborn as sns
|
||||
import torch
|
||||
import yaml
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from scipy.signal import butter, filtfilt
|
||||
|
||||
from utils.general import xywh2xyxy, xyxy2xywh
|
||||
from utils.metrics import fitness
|
||||
@ -26,12 +25,25 @@ matplotlib.rc('font', **{'size': 11})
|
||||
matplotlib.use('Agg') # for writing to files only
|
||||
|
||||
|
||||
def color_list():
|
||||
# Return first 10 plt colors as (r,g,b) https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb
|
||||
def hex2rgb(h):
|
||||
class Colors:
|
||||
# Ultralytics color palette https://ultralytics.com/
|
||||
def __init__(self):
|
||||
# hex = matplotlib.colors.TABLEAU_COLORS.values()
|
||||
hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
|
||||
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
|
||||
self.palette = [self.hex2rgb('#' + c) for c in hex]
|
||||
self.n = len(self.palette)
|
||||
|
||||
def __call__(self, i, bgr=False):
|
||||
c = self.palette[int(i) % self.n]
|
||||
return (c[2], c[1], c[0]) if bgr else c
|
||||
|
||||
@staticmethod
|
||||
def hex2rgb(h): # rgb order (PIL)
|
||||
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
|
||||
|
||||
return [hex2rgb(h) for h in matplotlib.colors.TABLEAU_COLORS.values()] # or BASE_ (8), CSS4_ (148), XKCD_ (949)
|
||||
|
||||
colors = Colors() # create instance for 'from utils.plots import colors'
|
||||
|
||||
|
||||
def hist2d(x, y, n=100):
|
||||
@ -44,6 +56,8 @@ def hist2d(x, y, n=100):
|
||||
|
||||
|
||||
def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
||||
from scipy.signal import butter, filtfilt
|
||||
|
||||
# https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy
|
||||
def butter_lowpass(cutoff, fs, order):
|
||||
nyq = 0.5 * fs
|
||||
@ -54,32 +68,32 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
||||
return filtfilt(b, a, data) # forward-backward filter
|
||||
|
||||
|
||||
def plot_one_box(x, img, color=None, label=None, line_thickness=3):
|
||||
# Plots one bounding box on image img
|
||||
tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
|
||||
color = color or [random.randint(0, 255) for _ in range(3)]
|
||||
def plot_one_box(x, im, color=(128, 128, 128), label=None, line_thickness=3):
|
||||
# Plots one bounding box on image 'im' using OpenCV
|
||||
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
|
||||
tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 # line/font thickness
|
||||
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
|
||||
cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
|
||||
cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
|
||||
if label:
|
||||
tf = max(tl - 1, 1) # font thickness
|
||||
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
|
||||
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
||||
cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
|
||||
cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
||||
cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled
|
||||
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
||||
|
||||
|
||||
def plot_one_box_PIL(box, img, color=None, label=None, line_thickness=None):
|
||||
img = Image.fromarray(img)
|
||||
draw = ImageDraw.Draw(img)
|
||||
line_thickness = line_thickness or max(int(min(img.size) / 200), 2)
|
||||
draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot
|
||||
def plot_one_box_PIL(box, im, color=(128, 128, 128), label=None, line_thickness=None):
|
||||
# Plots one bounding box on image 'im' using PIL
|
||||
im = Image.fromarray(im)
|
||||
draw = ImageDraw.Draw(im)
|
||||
line_thickness = line_thickness or max(int(min(im.size) / 200), 2)
|
||||
draw.rectangle(box, width=line_thickness, outline=color) # plot
|
||||
if label:
|
||||
fontsize = max(round(max(img.size) / 40), 12)
|
||||
font = ImageFont.truetype("Arial.ttf", fontsize)
|
||||
font = ImageFont.truetype("Arial.ttf", size=max(round(max(im.size) / 40), 12))
|
||||
txt_width, txt_height = font.getsize(label)
|
||||
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
|
||||
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=color)
|
||||
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
|
||||
return np.asarray(img)
|
||||
return np.asarray(im)
|
||||
|
||||
|
||||
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
|
||||
@ -135,7 +149,6 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max
|
||||
h = math.ceil(scale_factor * h)
|
||||
w = math.ceil(scale_factor * w)
|
||||
|
||||
colors = color_list() # list of colors
|
||||
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
||||
for i, img in enumerate(images):
|
||||
if i == max_subplots: # if last batch has fewer images than we expect
|
||||
@ -166,7 +179,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max
|
||||
boxes[[1, 3]] += block_y
|
||||
for j, box in enumerate(boxes.T):
|
||||
cls = int(classes[j])
|
||||
color = colors[cls % len(colors)]
|
||||
color = colors(cls)
|
||||
cls = names[cls] if names else cls
|
||||
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
||||
label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j])
|
||||
@ -274,7 +287,6 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
||||
print('Plotting labels... ')
|
||||
c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
|
||||
nc = int(c.max() + 1) # number of classes
|
||||
colors = color_list()
|
||||
x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
|
||||
|
||||
# seaborn correlogram
|
||||
@ -285,7 +297,8 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
||||
# matplotlib labels
|
||||
matplotlib.use('svg') # faster
|
||||
ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
|
||||
ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
|
||||
y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
|
||||
# [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # update colors bug #3195
|
||||
ax[0].set_ylabel('instances')
|
||||
if 0 < len(names) < 30:
|
||||
ax[0].set_xticks(range(len(names)))
|
||||
@ -300,7 +313,7 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
||||
labels[:, 1:] = xywh2xyxy(labels[:, 1:]) * 2000
|
||||
img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255)
|
||||
for cls, *box in labels[:1000]:
|
||||
ImageDraw.Draw(img).rectangle(box, width=1, outline=colors[int(cls) % 10]) # plot
|
||||
ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot
|
||||
ax[1].imshow(img)
|
||||
ax[1].axis('off')
|
||||
|
||||
@ -321,7 +334,7 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
||||
def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution()
|
||||
# Plot hyperparameter evolution results in evolve.txt
|
||||
with open(yaml_file) as f:
|
||||
hyp = yaml.load(f, Loader=yaml.SafeLoader)
|
||||
hyp = yaml.safe_load(f)
|
||||
x = np.loadtxt('evolve.txt', ndmin=2)
|
||||
f = fitness(x)
|
||||
# weights = (f - f.min()) ** 2 # for weighted results
|
||||
|
||||
@ -72,11 +72,12 @@ def select_device(device='', batch_size=None):
|
||||
|
||||
cuda = not cpu and torch.cuda.is_available()
|
||||
if cuda:
|
||||
n = torch.cuda.device_count()
|
||||
if n > 1 and batch_size: # check that batch_size is compatible with device_count
|
||||
devices = device.split(',') if device else range(torch.cuda.device_count()) # i.e. 0,1,6,7
|
||||
n = len(devices) # device count
|
||||
if n > 1 and batch_size: # check batch_size is divisible by device_count
|
||||
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
|
||||
space = ' ' * len(s)
|
||||
for i, d in enumerate(device.split(',') if device else range(n)):
|
||||
for i, d in enumerate(devices):
|
||||
p = torch.cuda.get_device_properties(i)
|
||||
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
|
||||
else:
|
||||
@ -133,9 +134,15 @@ def profile(x, ops, n=100, device=None):
|
||||
|
||||
|
||||
def is_parallel(model):
|
||||
# Returns True if model is of type DP or DDP
|
||||
return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
|
||||
|
||||
|
||||
def de_parallel(model):
|
||||
# De-parallelize a model: returns single-GPU model if model is of type DP or DDP
|
||||
return model.module if is_parallel(model) else model
|
||||
|
||||
|
||||
def intersect_dicts(da, db, exclude=()):
|
||||
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
|
||||
return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}
|
||||
|
||||
@ -9,7 +9,7 @@ WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
||||
|
||||
def create_dataset_artifact(opt):
|
||||
with open(opt.data) as f:
|
||||
data = yaml.load(f, Loader=yaml.SafeLoader) # data dict
|
||||
data = yaml.safe_load(f) # data dict
|
||||
logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation')
|
||||
|
||||
|
||||
@ -17,7 +17,7 @@ if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
|
||||
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
|
||||
parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
|
||||
parser.add_argument('--project', type=str, default='YOLOv3', help='name of W&B Project')
|
||||
opt = parser.parse_args()
|
||||
opt.resume = False # Explicitly disallow resume check for dataset upload job
|
||||
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
"""Utilities and tools for tracking runs with Weights & Biases."""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
@ -9,7 +10,7 @@ from tqdm import tqdm
|
||||
sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path
|
||||
from utils.datasets import LoadImagesAndLabels
|
||||
from utils.datasets import img2label_paths
|
||||
from utils.general import colorstr, xywh2xyxy, check_dataset
|
||||
from utils.general import colorstr, xywh2xyxy, check_dataset, check_file
|
||||
|
||||
try:
|
||||
import wandb
|
||||
@ -35,8 +36,9 @@ def get_run_info(run_path):
|
||||
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
|
||||
run_id = run_path.stem
|
||||
project = run_path.parent.stem
|
||||
entity = run_path.parent.parent.stem
|
||||
model_artifact_name = 'run_' + run_id + '_model'
|
||||
return run_id, project, model_artifact_name
|
||||
return entity, project, run_id, model_artifact_name
|
||||
|
||||
|
||||
def check_wandb_resume(opt):
|
||||
@ -44,9 +46,9 @@ def check_wandb_resume(opt):
|
||||
if isinstance(opt.resume, str):
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
if opt.global_rank not in [-1, 0]: # For resuming DDP runs
|
||||
run_id, project, model_artifact_name = get_run_info(opt.resume)
|
||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||
api = wandb.Api()
|
||||
artifact = api.artifact(project + '/' + model_artifact_name + ':latest')
|
||||
artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest')
|
||||
modeldir = artifact.download()
|
||||
opt.weights = str(Path(modeldir) / "last.pt")
|
||||
return True
|
||||
@ -54,8 +56,8 @@ def check_wandb_resume(opt):
|
||||
|
||||
|
||||
def process_wandb_config_ddp_mode(opt):
|
||||
with open(opt.data) as f:
|
||||
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
|
||||
with open(check_file(opt.data)) as f:
|
||||
data_dict = yaml.safe_load(f) # data dict
|
||||
train_dir, val_dir = None, None
|
||||
if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
|
||||
api = wandb.Api()
|
||||
@ -73,11 +75,23 @@ def process_wandb_config_ddp_mode(opt):
|
||||
if train_dir or val_dir:
|
||||
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
|
||||
with open(ddp_data_path, 'w') as f:
|
||||
yaml.dump(data_dict, f)
|
||||
yaml.safe_dump(data_dict, f)
|
||||
opt.data = ddp_data_path
|
||||
|
||||
|
||||
class WandbLogger():
|
||||
"""Log training runs, datasets, models, and predictions to Weights & Biases.
|
||||
|
||||
This logger sends information to W&B at wandb.ai. By default, this information
|
||||
includes hyperparameters, system configuration and metrics, model metrics,
|
||||
and basic data metrics and analyses.
|
||||
|
||||
By providing additional command line arguments to train.py, datasets,
|
||||
models and predictions can also be logged.
|
||||
|
||||
For more on how this logger is used, see the Weights & Biases documentation:
|
||||
https://docs.wandb.com/guides/integrations/yolov5
|
||||
"""
|
||||
def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
|
||||
# Pre-training routine --
|
||||
self.job_type = job_type
|
||||
@ -85,16 +99,17 @@ class WandbLogger():
|
||||
# It's more elegant to stick to 1 wandb.init call, but useful config data is overwritten in the WandbLogger's wandb.init call
|
||||
if isinstance(opt.resume, str): # checks resume from artifact
|
||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
run_id, project, model_artifact_name = get_run_info(opt.resume)
|
||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
|
||||
assert wandb, 'install wandb to resume wandb runs'
|
||||
# Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
|
||||
self.wandb_run = wandb.init(id=run_id, project=project, resume='allow')
|
||||
self.wandb_run = wandb.init(id=run_id, project=project, entity=entity, resume='allow')
|
||||
opt.resume = model_artifact_name
|
||||
elif self.wandb:
|
||||
self.wandb_run = wandb.init(config=opt,
|
||||
resume="allow",
|
||||
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
|
||||
project='YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem,
|
||||
entity=opt.entity,
|
||||
name=name,
|
||||
job_type=job_type,
|
||||
id=run_id) if not wandb.run else wandb.run
|
||||
@ -110,17 +125,17 @@ class WandbLogger():
|
||||
self.data_dict = self.check_and_upload_dataset(opt)
|
||||
else:
|
||||
prefix = colorstr('wandb: ')
|
||||
print(f"{prefix}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
|
||||
print(f"{prefix}Install Weights & Biases for YOLOv3 logging with 'pip install wandb' (recommended)")
|
||||
|
||||
def check_and_upload_dataset(self, opt):
|
||||
assert wandb, 'Install wandb to upload dataset'
|
||||
check_dataset(self.data_dict)
|
||||
config_path = self.log_dataset_artifact(opt.data,
|
||||
config_path = self.log_dataset_artifact(check_file(opt.data),
|
||||
opt.single_cls,
|
||||
'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
|
||||
'YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem)
|
||||
print("Created dataset config file ", config_path)
|
||||
with open(config_path) as f:
|
||||
wandb_data_dict = yaml.load(f, Loader=yaml.SafeLoader)
|
||||
wandb_data_dict = yaml.safe_load(f)
|
||||
return wandb_data_dict
|
||||
|
||||
def setup_training(self, opt, data_dict):
|
||||
@ -158,7 +173,8 @@ class WandbLogger():
|
||||
|
||||
def download_dataset_artifact(self, path, alias):
|
||||
if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
|
||||
dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
|
||||
artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
|
||||
dataset_artifact = wandb.use_artifact(artifact_path.as_posix())
|
||||
assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
|
||||
datadir = dataset_artifact.download()
|
||||
return datadir, dataset_artifact
|
||||
@ -171,8 +187,8 @@ class WandbLogger():
|
||||
modeldir = model_artifact.download()
|
||||
epochs_trained = model_artifact.metadata.get('epochs_trained')
|
||||
total_epochs = model_artifact.metadata.get('total_epochs')
|
||||
assert epochs_trained < total_epochs, 'training to %g epochs is finished, nothing to resume.' % (
|
||||
total_epochs)
|
||||
is_finished = total_epochs is None
|
||||
assert not is_finished, 'training is finished, can only resume incomplete runs.'
|
||||
return modeldir, model_artifact
|
||||
return None, None
|
||||
|
||||
@ -187,18 +203,18 @@ class WandbLogger():
|
||||
})
|
||||
model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
|
||||
wandb.log_artifact(model_artifact,
|
||||
aliases=['latest', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
|
||||
aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
|
||||
print("Saving model artifact on epoch ", epoch + 1)
|
||||
|
||||
def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
|
||||
with open(data_file) as f:
|
||||
data = yaml.load(f, Loader=yaml.SafeLoader) # data dict
|
||||
data = yaml.safe_load(f) # data dict
|
||||
nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
|
||||
names = {k: v for k, v in enumerate(names)} # to index dictionary
|
||||
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||
data['train']), names, name='train') if data.get('train') else None
|
||||
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None
|
||||
self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||
data['val']), names, name='val') if data.get('val') else None
|
||||
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
|
||||
if data.get('train'):
|
||||
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
|
||||
if data.get('val'):
|
||||
@ -206,7 +222,7 @@ class WandbLogger():
|
||||
path = data_file if overwrite_config else '_wandb.'.join(data_file.rsplit('.', 1)) # updated data.yaml path
|
||||
data.pop('download', None)
|
||||
with open(path, 'w') as f:
|
||||
yaml.dump(data, f)
|
||||
yaml.safe_dump(data, f)
|
||||
|
||||
if self.job_type == 'Training': # builds correct artifact pipeline graph
|
||||
self.wandb_run.use_artifact(self.val_artifact)
|
||||
@ -243,16 +259,12 @@ class WandbLogger():
|
||||
table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
|
||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
|
||||
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
|
||||
height, width = shapes[0]
|
||||
labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4))) * torch.Tensor([width, height, width, height])
|
||||
box_data, img_classes = [], {}
|
||||
for cls, *xyxy in labels[:, 1:].tolist():
|
||||
for cls, *xywh in labels[:, 1:].tolist():
|
||||
cls = int(cls)
|
||||
box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]},
|
||||
"class_id": cls,
|
||||
"box_caption": "%s" % (class_to_id[cls]),
|
||||
"scores": {"acc": 1},
|
||||
"domain": "pixel"})
|
||||
"box_caption": "%s" % (class_to_id[cls])})
|
||||
img_classes[cls] = class_to_id[cls]
|
||||
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
|
||||
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes),
|
||||
@ -294,7 +306,7 @@ class WandbLogger():
|
||||
if self.result_artifact:
|
||||
train_results = wandb.JoinedTable(self.val_table, self.result_table, "id")
|
||||
self.result_artifact.add(train_results, 'result')
|
||||
wandb.log_artifact(self.result_artifact, aliases=['latest', 'epoch ' + str(self.current_epoch),
|
||||
wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch),
|
||||
('best' if best_result else '')])
|
||||
self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
|
||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user