Upload 161 files

ultralytics/.pre-commit-config.yaml

@@ -0,0 +1,73 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Pre-commit hooks. For more information see https://github.com/pre-commit/pre-commit-hooks/blob/main/README.md
+
+exclude: 'docs/'
+# Define bot property if installed via https://github.com/marketplace/pre-commit-ci
+ci:
+  autofix_prs: true
+  autoupdate_commit_msg: '[pre-commit.ci] pre-commit suggestions'
+  autoupdate_schedule: monthly
+  # submodules: true
+
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: check-case-conflict
+      # - id: check-yaml
+      - id: check-docstring-first
+      - id: double-quote-string-fixer
+      - id: detect-private-key
+
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.4.0
+    hooks:
+      - id: pyupgrade
+        name: Upgrade code
+
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        name: Sort imports
+
+  - repo: https://github.com/google/yapf
+    rev: v0.33.0
+    hooks:
+      - id: yapf
+        name: YAPF formatting
+
+  - repo: https://github.com/executablebooks/mdformat
+    rev: 0.7.16
+    hooks:
+      - id: mdformat
+        name: MD formatting
+        additional_dependencies:
+          - mdformat-gfm
+          - mdformat-black
+        # exclude: "README.md|README.zh-CN.md|CONTRIBUTING.md"
+
+  - repo: https://github.com/PyCQA/flake8
+    rev: 6.0.0
+    hooks:
+      - id: flake8
+        name: PEP8
+
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.2.4
+    hooks:
+      - id: codespell
+        args:
+          - --ignore-words-list=crate,nd,strack,dota
+
+#  - repo: https://github.com/asottile/yesqa
+#    rev: v1.4.0
+#    hooks:
+#      - id: yesqa
+
+#  - repo: https://github.com/asottile/dead
+#    rev: v1.5.0
+#    hooks:
+#    -   id: dead

ultralytics/__init__.py

@@ -0,0 +1,12 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+__version__ = '8.0.120'
+
+from ultralytics.hub import start
+from ultralytics.vit.rtdetr import RTDETR
+from ultralytics.vit.sam import SAM
+from ultralytics.yolo.engine.model import YOLO
+from ultralytics.yolo.nas import NAS
+from ultralytics.yolo.utils.checks import check_yolo as checks
+
+__all__ = '__version__', 'YOLO', 'NAS', 'SAM', 'RTDETR', 'checks', 'start'  # allow simpler import

ultralytics/datasets/Argoverse.yaml

@@ -0,0 +1,73 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI
+# Example usage: yolo train data=Argoverse.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── Argoverse  ← downloads here (31.3 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/Argoverse  # dataset root dir
+train: Argoverse-1.1/images/train/  # train images (relative to 'path') 39384 images
+val: Argoverse-1.1/images/val/  # val images (relative to 'path') 15062 images
+test: Argoverse-1.1/images/test/  # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: bus
+  5: truck
+  6: traffic_light
+  7: stop_sign
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import json
+  from tqdm import tqdm
+  from ultralytics.yolo.utils.downloads import download
+  from pathlib import Path
+
+  def argoverse2yolo(set):
+      labels = {}
+      a = json.load(open(set, "rb"))
+      for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv5 format..."):
+          img_id = annot['image_id']
+          img_name = a['images'][img_id]['name']
+          img_label_name = f'{img_name[:-3]}txt'
+
+          cls = annot['category_id']  # instance class id
+          x_center, y_center, width, height = annot['bbox']
+          x_center = (x_center + width / 2) / 1920.0  # offset and scale
+          y_center = (y_center + height / 2) / 1200.0  # offset and scale
+          width /= 1920.0  # scale
+          height /= 1200.0  # scale
+
+          img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']]
+          if not img_dir.exists():
+              img_dir.mkdir(parents=True, exist_ok=True)
+
+          k = str(img_dir / img_label_name)
+          if k not in labels:
+              labels[k] = []
+          labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n")
+
+      for k in labels:
+          with open(k, "w") as f:
+              f.writelines(labels[k])
+
+
+  # Download
+  dir = Path(yaml['path'])  # dataset root dir
+  urls = ['https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip']
+  download(urls, dir=dir)
+
+  # Convert
+  annotations_dir = 'Argoverse-HD/annotations/'
+  (dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images')  # rename 'tracking' to 'images'
+  for d in "train.json", "val.json":
+      argoverse2yolo(dir / annotations_dir / d)  # convert VisDrone annotations to YOLO labels

ultralytics/datasets/GlobalWheat2020.yaml

@@ -0,0 +1,54 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Global Wheat 2020 dataset http://www.global-wheat.com/ by University of Saskatchewan
+# Example usage: yolo train data=GlobalWheat2020.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── GlobalWheat2020  ← downloads here (7.0 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/GlobalWheat2020  # dataset root dir
+train: # train images (relative to 'path') 3422 images
+  - images/arvalis_1
+  - images/arvalis_2
+  - images/arvalis_3
+  - images/ethz_1
+  - images/rres_1
+  - images/inrae_1
+  - images/usask_1
+val: # val images (relative to 'path') 748 images (WARNING: train set contains ethz_1)
+  - images/ethz_1
+test: # test images (optional) 1276 images
+  - images/utokyo_1
+  - images/utokyo_2
+  - images/nau_1
+  - images/uq_1
+
+# Classes
+names:
+  0: wheat_head
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  from ultralytics.yolo.utils.downloads import download
+  from pathlib import Path
+
+  # Download
+  dir = Path(yaml['path'])  # dataset root dir
+  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 / 'global-wheat-codalab-official' / p).rename(dir / 'images' / p)  # move to /images
+      f = (dir / 'global-wheat-codalab-official' / p).with_suffix('.json')  # json file
+      if f.exists():
+          f.rename((dir / 'annotations' / p).with_suffix('.json'))  # move to /annotations

ultralytics/datasets/ImageNet.yaml

@@ -0,0 +1,2025 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
+# Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
+# Example usage: yolo train task=classify data=imagenet
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── imagenet  ← downloads here (144 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/imagenet  # dataset root dir
+train: train  # train images (relative to 'path') 1281167 images
+val: val  # val images (relative to 'path') 50000 images
+test:  # test images (optional)
+
+# Classes
+names:
+  0: tench
+  1: goldfish
+  2: great white shark
+  3: tiger shark
+  4: hammerhead shark
+  5: electric ray
+  6: stingray
+  7: cock
+  8: hen
+  9: ostrich
+  10: brambling
+  11: goldfinch
+  12: house finch
+  13: junco
+  14: indigo bunting
+  15: American robin
+  16: bulbul
+  17: jay
+  18: magpie
+  19: chickadee
+  20: American dipper
+  21: kite
+  22: bald eagle
+  23: vulture
+  24: great grey owl
+  25: fire salamander
+  26: smooth newt
+  27: newt
+  28: spotted salamander
+  29: axolotl
+  30: American bullfrog
+  31: tree frog
+  32: tailed frog
+  33: loggerhead sea turtle
+  34: leatherback sea turtle
+  35: mud turtle
+  36: terrapin
+  37: box turtle
+  38: banded gecko
+  39: green iguana
+  40: Carolina anole
+  41: desert grassland whiptail lizard
+  42: agama
+  43: frilled-necked lizard
+  44: alligator lizard
+  45: Gila monster
+  46: European green lizard
+  47: chameleon
+  48: Komodo dragon
+  49: Nile crocodile
+  50: American alligator
+  51: triceratops
+  52: worm snake
+  53: ring-necked snake
+  54: eastern hog-nosed snake
+  55: smooth green snake
+  56: kingsnake
+  57: garter snake
+  58: water snake
+  59: vine snake
+  60: night snake
+  61: boa constrictor
+  62: African rock python
+  63: Indian cobra
+  64: green mamba
+  65: sea snake
+  66: Saharan horned viper
+  67: eastern diamondback rattlesnake
+  68: sidewinder
+  69: trilobite
+  70: harvestman
+  71: scorpion
+  72: yellow garden spider
+  73: barn spider
+  74: European garden spider
+  75: southern black widow
+  76: tarantula
+  77: wolf spider
+  78: tick
+  79: centipede
+  80: black grouse
+  81: ptarmigan
+  82: ruffed grouse
+  83: prairie grouse
+  84: peacock
+  85: quail
+  86: partridge
+  87: grey parrot
+  88: macaw
+  89: sulphur-crested cockatoo
+  90: lorikeet
+  91: coucal
+  92: bee eater
+  93: hornbill
+  94: hummingbird
+  95: jacamar
+  96: toucan
+  97: duck
+  98: red-breasted merganser
+  99: goose
+  100: black swan
+  101: tusker
+  102: echidna
+  103: platypus
+  104: wallaby
+  105: koala
+  106: wombat
+  107: jellyfish
+  108: sea anemone
+  109: brain coral
+  110: flatworm
+  111: nematode
+  112: conch
+  113: snail
+  114: slug
+  115: sea slug
+  116: chiton
+  117: chambered nautilus
+  118: Dungeness crab
+  119: rock crab
+  120: fiddler crab
+  121: red king crab
+  122: American lobster
+  123: spiny lobster
+  124: crayfish
+  125: hermit crab
+  126: isopod
+  127: white stork
+  128: black stork
+  129: spoonbill
+  130: flamingo
+  131: little blue heron
+  132: great egret
+  133: bittern
+  134: crane (bird)
+  135: limpkin
+  136: common gallinule
+  137: American coot
+  138: bustard
+  139: ruddy turnstone
+  140: dunlin
+  141: common redshank
+  142: dowitcher
+  143: oystercatcher
+  144: pelican
+  145: king penguin
+  146: albatross
+  147: grey whale
+  148: killer whale
+  149: dugong
+  150: sea lion
+  151: Chihuahua
+  152: Japanese Chin
+  153: Maltese
+  154: Pekingese
+  155: Shih Tzu
+  156: King Charles Spaniel
+  157: Papillon
+  158: toy terrier
+  159: Rhodesian Ridgeback
+  160: Afghan Hound
+  161: Basset Hound
+  162: Beagle
+  163: Bloodhound
+  164: Bluetick Coonhound
+  165: Black and Tan Coonhound
+  166: Treeing Walker Coonhound
+  167: English foxhound
+  168: Redbone Coonhound
+  169: borzoi
+  170: Irish Wolfhound
+  171: Italian Greyhound
+  172: Whippet
+  173: Ibizan Hound
+  174: Norwegian Elkhound
+  175: Otterhound
+  176: Saluki
+  177: Scottish Deerhound
+  178: Weimaraner
+  179: Staffordshire Bull Terrier
+  180: American Staffordshire Terrier
+  181: Bedlington Terrier
+  182: Border Terrier
+  183: Kerry Blue Terrier
+  184: Irish Terrier
+  185: Norfolk Terrier
+  186: Norwich Terrier
+  187: Yorkshire Terrier
+  188: Wire Fox Terrier
+  189: Lakeland Terrier
+  190: Sealyham Terrier
+  191: Airedale Terrier
+  192: Cairn Terrier
+  193: Australian Terrier
+  194: Dandie Dinmont Terrier
+  195: Boston Terrier
+  196: Miniature Schnauzer
+  197: Giant Schnauzer
+  198: Standard Schnauzer
+  199: Scottish Terrier
+  200: Tibetan Terrier
+  201: Australian Silky Terrier
+  202: Soft-coated Wheaten Terrier
+  203: West Highland White Terrier
+  204: Lhasa Apso
+  205: Flat-Coated Retriever
+  206: Curly-coated Retriever
+  207: Golden Retriever
+  208: Labrador Retriever
+  209: Chesapeake Bay Retriever
+  210: German Shorthaired Pointer
+  211: Vizsla
+  212: English Setter
+  213: Irish Setter
+  214: Gordon Setter
+  215: Brittany
+  216: Clumber Spaniel
+  217: English Springer Spaniel
+  218: Welsh Springer Spaniel
+  219: Cocker Spaniels
+  220: Sussex Spaniel
+  221: Irish Water Spaniel
+  222: Kuvasz
+  223: Schipperke
+  224: Groenendael
+  225: Malinois
+  226: Briard
+  227: Australian Kelpie
+  228: Komondor
+  229: Old English Sheepdog
+  230: Shetland Sheepdog
+  231: collie
+  232: Border Collie
+  233: Bouvier des Flandres
+  234: Rottweiler
+  235: German Shepherd Dog
+  236: Dobermann
+  237: Miniature Pinscher
+  238: Greater Swiss Mountain Dog
+  239: Bernese Mountain Dog
+  240: Appenzeller Sennenhund
+  241: Entlebucher Sennenhund
+  242: Boxer
+  243: Bullmastiff
+  244: Tibetan Mastiff
+  245: French Bulldog
+  246: Great Dane
+  247: St. Bernard
+  248: husky
+  249: Alaskan Malamute
+  250: Siberian Husky
+  251: Dalmatian
+  252: Affenpinscher
+  253: Basenji
+  254: pug
+  255: Leonberger
+  256: Newfoundland
+  257: Pyrenean Mountain Dog
+  258: Samoyed
+  259: Pomeranian
+  260: Chow Chow
+  261: Keeshond
+  262: Griffon Bruxellois
+  263: Pembroke Welsh Corgi
+  264: Cardigan Welsh Corgi
+  265: Toy Poodle
+  266: Miniature Poodle
+  267: Standard Poodle
+  268: Mexican hairless dog
+  269: grey wolf
+  270: Alaskan tundra wolf
+  271: red wolf
+  272: coyote
+  273: dingo
+  274: dhole
+  275: African wild dog
+  276: hyena
+  277: red fox
+  278: kit fox
+  279: Arctic fox
+  280: grey fox
+  281: tabby cat
+  282: tiger cat
+  283: Persian cat
+  284: Siamese cat
+  285: Egyptian Mau
+  286: cougar
+  287: lynx
+  288: leopard
+  289: snow leopard
+  290: jaguar
+  291: lion
+  292: tiger
+  293: cheetah
+  294: brown bear
+  295: American black bear
+  296: polar bear
+  297: sloth bear
+  298: mongoose
+  299: meerkat
+  300: tiger beetle
+  301: ladybug
+  302: ground beetle
+  303: longhorn beetle
+  304: leaf beetle
+  305: dung beetle
+  306: rhinoceros beetle
+  307: weevil
+  308: fly
+  309: bee
+  310: ant
+  311: grasshopper
+  312: cricket
+  313: stick insect
+  314: cockroach
+  315: mantis
+  316: cicada
+  317: leafhopper
+  318: lacewing
+  319: dragonfly
+  320: damselfly
+  321: red admiral
+  322: ringlet
+  323: monarch butterfly
+  324: small white
+  325: sulphur butterfly
+  326: gossamer-winged butterfly
+  327: starfish
+  328: sea urchin
+  329: sea cucumber
+  330: cottontail rabbit
+  331: hare
+  332: Angora rabbit
+  333: hamster
+  334: porcupine
+  335: fox squirrel
+  336: marmot
+  337: beaver
+  338: guinea pig
+  339: common sorrel
+  340: zebra
+  341: pig
+  342: wild boar
+  343: warthog
+  344: hippopotamus
+  345: ox
+  346: water buffalo
+  347: bison
+  348: ram
+  349: bighorn sheep
+  350: Alpine ibex
+  351: hartebeest
+  352: impala
+  353: gazelle
+  354: dromedary
+  355: llama
+  356: weasel
+  357: mink
+  358: European polecat
+  359: black-footed ferret
+  360: otter
+  361: skunk
+  362: badger
+  363: armadillo
+  364: three-toed sloth
+  365: orangutan
+  366: gorilla
+  367: chimpanzee
+  368: gibbon
+  369: siamang
+  370: guenon
+  371: patas monkey
+  372: baboon
+  373: macaque
+  374: langur
+  375: black-and-white colobus
+  376: proboscis monkey
+  377: marmoset
+  378: white-headed capuchin
+  379: howler monkey
+  380: titi
+  381: Geoffroy's spider monkey
+  382: common squirrel monkey
+  383: ring-tailed lemur
+  384: indri
+  385: Asian elephant
+  386: African bush elephant
+  387: red panda
+  388: giant panda
+  389: snoek
+  390: eel
+  391: coho salmon
+  392: rock beauty
+  393: clownfish
+  394: sturgeon
+  395: garfish
+  396: lionfish
+  397: pufferfish
+  398: abacus
+  399: abaya
+  400: academic gown
+  401: accordion
+  402: acoustic guitar
+  403: aircraft carrier
+  404: airliner
+  405: airship
+  406: altar
+  407: ambulance
+  408: amphibious vehicle
+  409: analog clock
+  410: apiary
+  411: apron
+  412: waste container
+  413: assault rifle
+  414: backpack
+  415: bakery
+  416: balance beam
+  417: balloon
+  418: ballpoint pen
+  419: Band-Aid
+  420: banjo
+  421: baluster
+  422: barbell
+  423: barber chair
+  424: barbershop
+  425: barn
+  426: barometer
+  427: barrel
+  428: wheelbarrow
+  429: baseball
+  430: basketball
+  431: bassinet
+  432: bassoon
+  433: swimming cap
+  434: bath towel
+  435: bathtub
+  436: station wagon
+  437: lighthouse
+  438: beaker
+  439: military cap
+  440: beer bottle
+  441: beer glass
+  442: bell-cot
+  443: bib
+  444: tandem bicycle
+  445: bikini
+  446: ring binder
+  447: binoculars
+  448: birdhouse
+  449: boathouse
+  450: bobsleigh
+  451: bolo tie
+  452: poke bonnet
+  453: bookcase
+  454: bookstore
+  455: bottle cap
+  456: bow
+  457: bow tie
+  458: brass
+  459: bra
+  460: breakwater
+  461: breastplate
+  462: broom
+  463: bucket
+  464: buckle
+  465: bulletproof vest
+  466: high-speed train
+  467: butcher shop
+  468: taxicab
+  469: cauldron
+  470: candle
+  471: cannon
+  472: canoe
+  473: can opener
+  474: cardigan
+  475: car mirror
+  476: carousel
+  477: tool kit
+  478: carton
+  479: car wheel
+  480: automated teller machine
+  481: cassette
+  482: cassette player
+  483: castle
+  484: catamaran
+  485: CD player
+  486: cello
+  487: mobile phone
+  488: chain
+  489: chain-link fence
+  490: chain mail
+  491: chainsaw
+  492: chest
+  493: chiffonier
+  494: chime
+  495: china cabinet
+  496: Christmas stocking
+  497: church
+  498: movie theater
+  499: cleaver
+  500: cliff dwelling
+  501: cloak
+  502: clogs
+  503: cocktail shaker
+  504: coffee mug
+  505: coffeemaker
+  506: coil
+  507: combination lock
+  508: computer keyboard
+  509: confectionery store
+  510: container ship
+  511: convertible
+  512: corkscrew
+  513: cornet
+  514: cowboy boot
+  515: cowboy hat
+  516: cradle
+  517: crane (machine)
+  518: crash helmet
+  519: crate
+  520: infant bed
+  521: Crock Pot
+  522: croquet ball
+  523: crutch
+  524: cuirass
+  525: dam
+  526: desk
+  527: desktop computer
+  528: rotary dial telephone
+  529: diaper
+  530: digital clock
+  531: digital watch
+  532: dining table
+  533: dishcloth
+  534: dishwasher
+  535: disc brake
+  536: dock
+  537: dog sled
+  538: dome
+  539: doormat
+  540: drilling rig
+  541: drum
+  542: drumstick
+  543: dumbbell
+  544: Dutch oven
+  545: electric fan
+  546: electric guitar
+  547: electric locomotive
+  548: entertainment center
+  549: envelope
+  550: espresso machine
+  551: face powder
+  552: feather boa
+  553: filing cabinet
+  554: fireboat
+  555: fire engine
+  556: fire screen sheet
+  557: flagpole
+  558: flute
+  559: folding chair
+  560: football helmet
+  561: forklift
+  562: fountain
+  563: fountain pen
+  564: four-poster bed
+  565: freight car
+  566: French horn
+  567: frying pan
+  568: fur coat
+  569: garbage truck
+  570: gas mask
+  571: gas pump
+  572: goblet
+  573: go-kart
+  574: golf ball
+  575: golf cart
+  576: gondola
+  577: gong
+  578: gown
+  579: grand piano
+  580: greenhouse
+  581: grille
+  582: grocery store
+  583: guillotine
+  584: barrette
+  585: hair spray
+  586: half-track
+  587: hammer
+  588: hamper
+  589: hair dryer
+  590: hand-held computer
+  591: handkerchief
+  592: hard disk drive
+  593: harmonica
+  594: harp
+  595: harvester
+  596: hatchet
+  597: holster
+  598: home theater
+  599: honeycomb
+  600: hook
+  601: hoop skirt
+  602: horizontal bar
+  603: horse-drawn vehicle
+  604: hourglass
+  605: iPod
+  606: clothes iron
+  607: jack-o'-lantern
+  608: jeans
+  609: jeep
+  610: T-shirt
+  611: jigsaw puzzle
+  612: pulled rickshaw
+  613: joystick
+  614: kimono
+  615: knee pad
+  616: knot
+  617: lab coat
+  618: ladle
+  619: lampshade
+  620: laptop computer
+  621: lawn mower
+  622: lens cap
+  623: paper knife
+  624: library
+  625: lifeboat
+  626: lighter
+  627: limousine
+  628: ocean liner
+  629: lipstick
+  630: slip-on shoe
+  631: lotion
+  632: speaker
+  633: loupe
+  634: sawmill
+  635: magnetic compass
+  636: mail bag
+  637: mailbox
+  638: tights
+  639: tank suit
+  640: manhole cover
+  641: maraca
+  642: marimba
+  643: mask
+  644: match
+  645: maypole
+  646: maze
+  647: measuring cup
+  648: medicine chest
+  649: megalith
+  650: microphone
+  651: microwave oven
+  652: military uniform
+  653: milk can
+  654: minibus
+  655: miniskirt
+  656: minivan
+  657: missile
+  658: mitten
+  659: mixing bowl
+  660: mobile home
+  661: Model T
+  662: modem
+  663: monastery
+  664: monitor
+  665: moped
+  666: mortar
+  667: square academic cap
+  668: mosque
+  669: mosquito net
+  670: scooter
+  671: mountain bike
+  672: tent
+  673: computer mouse
+  674: mousetrap
+  675: moving van
+  676: muzzle
+  677: nail
+  678: neck brace
+  679: necklace
+  680: nipple
+  681: notebook computer
+  682: obelisk
+  683: oboe
+  684: ocarina
+  685: odometer
+  686: oil filter
+  687: organ
+  688: oscilloscope
+  689: overskirt
+  690: bullock cart
+  691: oxygen mask
+  692: packet
+  693: paddle
+  694: paddle wheel
+  695: padlock
+  696: paintbrush
+  697: pajamas
+  698: palace
+  699: pan flute
+  700: paper towel
+  701: parachute
+  702: parallel bars
+  703: park bench
+  704: parking meter
+  705: passenger car
+  706: patio
+  707: payphone
+  708: pedestal
+  709: pencil case
+  710: pencil sharpener
+  711: perfume
+  712: Petri dish
+  713: photocopier
+  714: plectrum
+  715: Pickelhaube
+  716: picket fence
+  717: pickup truck
+  718: pier
+  719: piggy bank
+  720: pill bottle
+  721: pillow
+  722: ping-pong ball
+  723: pinwheel
+  724: pirate ship
+  725: pitcher
+  726: hand plane
+  727: planetarium
+  728: plastic bag
+  729: plate rack
+  730: plow
+  731: plunger
+  732: Polaroid camera
+  733: pole
+  734: police van
+  735: poncho
+  736: billiard table
+  737: soda bottle
+  738: pot
+  739: potter's wheel
+  740: power drill
+  741: prayer rug
+  742: printer
+  743: prison
+  744: projectile
+  745: projector
+  746: hockey puck
+  747: punching bag
+  748: purse
+  749: quill
+  750: quilt
+  751: race car
+  752: racket
+  753: radiator
+  754: radio
+  755: radio telescope
+  756: rain barrel
+  757: recreational vehicle
+  758: reel
+  759: reflex camera
+  760: refrigerator
+  761: remote control
+  762: restaurant
+  763: revolver
+  764: rifle
+  765: rocking chair
+  766: rotisserie
+  767: eraser
+  768: rugby ball
+  769: ruler
+  770: running shoe
+  771: safe
+  772: safety pin
+  773: salt shaker
+  774: sandal
+  775: sarong
+  776: saxophone
+  777: scabbard
+  778: weighing scale
+  779: school bus
+  780: schooner
+  781: scoreboard
+  782: CRT screen
+  783: screw
+  784: screwdriver
+  785: seat belt
+  786: sewing machine
+  787: shield
+  788: shoe store
+  789: shoji
+  790: shopping basket
+  791: shopping cart
+  792: shovel
+  793: shower cap
+  794: shower curtain
+  795: ski
+  796: ski mask
+  797: sleeping bag
+  798: slide rule
+  799: sliding door
+  800: slot machine
+  801: snorkel
+  802: snowmobile
+  803: snowplow
+  804: soap dispenser
+  805: soccer ball
+  806: sock
+  807: solar thermal collector
+  808: sombrero
+  809: soup bowl
+  810: space bar
+  811: space heater
+  812: space shuttle
+  813: spatula
+  814: motorboat
+  815: spider web
+  816: spindle
+  817: sports car
+  818: spotlight
+  819: stage
+  820: steam locomotive
+  821: through arch bridge
+  822: steel drum
+  823: stethoscope
+  824: scarf
+  825: stone wall
+  826: stopwatch
+  827: stove
+  828: strainer
+  829: tram
+  830: stretcher
+  831: couch
+  832: stupa
+  833: submarine
+  834: suit
+  835: sundial
+  836: sunglass
+  837: sunglasses
+  838: sunscreen
+  839: suspension bridge
+  840: mop
+  841: sweatshirt
+  842: swimsuit
+  843: swing
+  844: switch
+  845: syringe
+  846: table lamp
+  847: tank
+  848: tape player
+  849: teapot
+  850: teddy bear
+  851: television
+  852: tennis ball
+  853: thatched roof
+  854: front curtain
+  855: thimble
+  856: threshing machine
+  857: throne
+  858: tile roof
+  859: toaster
+  860: tobacco shop
+  861: toilet seat
+  862: torch
+  863: totem pole
+  864: tow truck
+  865: toy store
+  866: tractor
+  867: semi-trailer truck
+  868: tray
+  869: trench coat
+  870: tricycle
+  871: trimaran
+  872: tripod
+  873: triumphal arch
+  874: trolleybus
+  875: trombone
+  876: tub
+  877: turnstile
+  878: typewriter keyboard
+  879: umbrella
+  880: unicycle
+  881: upright piano
+  882: vacuum cleaner
+  883: vase
+  884: vault
+  885: velvet
+  886: vending machine
+  887: vestment
+  888: viaduct
+  889: violin
+  890: volleyball
+  891: waffle iron
+  892: wall clock
+  893: wallet
+  894: wardrobe
+  895: military aircraft
+  896: sink
+  897: washing machine
+  898: water bottle
+  899: water jug
+  900: water tower
+  901: whiskey jug
+  902: whistle
+  903: wig
+  904: window screen
+  905: window shade
+  906: Windsor tie
+  907: wine bottle
+  908: wing
+  909: wok
+  910: wooden spoon
+  911: wool
+  912: split-rail fence
+  913: shipwreck
+  914: yawl
+  915: yurt
+  916: website
+  917: comic book
+  918: crossword
+  919: traffic sign
+  920: traffic light
+  921: dust jacket
+  922: menu
+  923: plate
+  924: guacamole
+  925: consomme
+  926: hot pot
+  927: trifle
+  928: ice cream
+  929: ice pop
+  930: baguette
+  931: bagel
+  932: pretzel
+  933: cheeseburger
+  934: hot dog
+  935: mashed potato
+  936: cabbage
+  937: broccoli
+  938: cauliflower
+  939: zucchini
+  940: spaghetti squash
+  941: acorn squash
+  942: butternut squash
+  943: cucumber
+  944: artichoke
+  945: bell pepper
+  946: cardoon
+  947: mushroom
+  948: Granny Smith
+  949: strawberry
+  950: orange
+  951: lemon
+  952: fig
+  953: pineapple
+  954: banana
+  955: jackfruit
+  956: custard apple
+  957: pomegranate
+  958: hay
+  959: carbonara
+  960: chocolate syrup
+  961: dough
+  962: meatloaf
+  963: pizza
+  964: pot pie
+  965: burrito
+  966: red wine
+  967: espresso
+  968: cup
+  969: eggnog
+  970: alp
+  971: bubble
+  972: cliff
+  973: coral reef
+  974: geyser
+  975: lakeshore
+  976: promontory
+  977: shoal
+  978: seashore
+  979: valley
+  980: volcano
+  981: baseball player
+  982: bridegroom
+  983: scuba diver
+  984: rapeseed
+  985: daisy
+  986: yellow lady's slipper
+  987: corn
+  988: acorn
+  989: rose hip
+  990: horse chestnut seed
+  991: coral fungus
+  992: agaric
+  993: gyromitra
+  994: stinkhorn mushroom
+  995: earth star
+  996: hen-of-the-woods
+  997: bolete
+  998: ear
+  999: toilet paper
+
+# Imagenet class codes to human-readable names
+map:
+  n01440764: tench
+  n01443537: goldfish
+  n01484850: great_white_shark
+  n01491361: tiger_shark
+  n01494475: hammerhead
+  n01496331: electric_ray
+  n01498041: stingray
+  n01514668: cock
+  n01514859: hen
+  n01518878: ostrich
+  n01530575: brambling
+  n01531178: goldfinch
+  n01532829: house_finch
+  n01534433: junco
+  n01537544: indigo_bunting
+  n01558993: robin
+  n01560419: bulbul
+  n01580077: jay
+  n01582220: magpie
+  n01592084: chickadee
+  n01601694: water_ouzel
+  n01608432: kite
+  n01614925: bald_eagle
+  n01616318: vulture
+  n01622779: great_grey_owl
+  n01629819: European_fire_salamander
+  n01630670: common_newt
+  n01631663: eft
+  n01632458: spotted_salamander
+  n01632777: axolotl
+  n01641577: bullfrog
+  n01644373: tree_frog
+  n01644900: tailed_frog
+  n01664065: loggerhead
+  n01665541: leatherback_turtle
+  n01667114: mud_turtle
+  n01667778: terrapin
+  n01669191: box_turtle
+  n01675722: banded_gecko
+  n01677366: common_iguana
+  n01682714: American_chameleon
+  n01685808: whiptail
+  n01687978: agama
+  n01688243: frilled_lizard
+  n01689811: alligator_lizard
+  n01692333: Gila_monster
+  n01693334: green_lizard
+  n01694178: African_chameleon
+  n01695060: Komodo_dragon
+  n01697457: African_crocodile
+  n01698640: American_alligator
+  n01704323: triceratops
+  n01728572: thunder_snake
+  n01728920: ringneck_snake
+  n01729322: hognose_snake
+  n01729977: green_snake
+  n01734418: king_snake
+  n01735189: garter_snake
+  n01737021: water_snake
+  n01739381: vine_snake
+  n01740131: night_snake
+  n01742172: boa_constrictor
+  n01744401: rock_python
+  n01748264: Indian_cobra
+  n01749939: green_mamba
+  n01751748: sea_snake
+  n01753488: horned_viper
+  n01755581: diamondback
+  n01756291: sidewinder
+  n01768244: trilobite
+  n01770081: harvestman
+  n01770393: scorpion
+  n01773157: black_and_gold_garden_spider
+  n01773549: barn_spider
+  n01773797: garden_spider
+  n01774384: black_widow
+  n01774750: tarantula
+  n01775062: wolf_spider
+  n01776313: tick
+  n01784675: centipede
+  n01795545: black_grouse
+  n01796340: ptarmigan
+  n01797886: ruffed_grouse
+  n01798484: prairie_chicken
+  n01806143: peacock
+  n01806567: quail
+  n01807496: partridge
+  n01817953: African_grey
+  n01818515: macaw
+  n01819313: sulphur-crested_cockatoo
+  n01820546: lorikeet
+  n01824575: coucal
+  n01828970: bee_eater
+  n01829413: hornbill
+  n01833805: hummingbird
+  n01843065: jacamar
+  n01843383: toucan
+  n01847000: drake
+  n01855032: red-breasted_merganser
+  n01855672: goose
+  n01860187: black_swan
+  n01871265: tusker
+  n01872401: echidna
+  n01873310: platypus
+  n01877812: wallaby
+  n01882714: koala
+  n01883070: wombat
+  n01910747: jellyfish
+  n01914609: sea_anemone
+  n01917289: brain_coral
+  n01924916: flatworm
+  n01930112: nematode
+  n01943899: conch
+  n01944390: snail
+  n01945685: slug
+  n01950731: sea_slug
+  n01955084: chiton
+  n01968897: chambered_nautilus
+  n01978287: Dungeness_crab
+  n01978455: rock_crab
+  n01980166: fiddler_crab
+  n01981276: king_crab
+  n01983481: American_lobster
+  n01984695: spiny_lobster
+  n01985128: crayfish
+  n01986214: hermit_crab
+  n01990800: isopod
+  n02002556: white_stork
+  n02002724: black_stork
+  n02006656: spoonbill
+  n02007558: flamingo
+  n02009229: little_blue_heron
+  n02009912: American_egret
+  n02011460: bittern
+  n02012849: crane_(bird)
+  n02013706: limpkin
+  n02017213: European_gallinule
+  n02018207: American_coot
+  n02018795: bustard
+  n02025239: ruddy_turnstone
+  n02027492: red-backed_sandpiper
+  n02028035: redshank
+  n02033041: dowitcher
+  n02037110: oystercatcher
+  n02051845: pelican
+  n02056570: king_penguin
+  n02058221: albatross
+  n02066245: grey_whale
+  n02071294: killer_whale
+  n02074367: dugong
+  n02077923: sea_lion
+  n02085620: Chihuahua
+  n02085782: Japanese_spaniel
+  n02085936: Maltese_dog
+  n02086079: Pekinese
+  n02086240: Shih-Tzu
+  n02086646: Blenheim_spaniel
+  n02086910: papillon
+  n02087046: toy_terrier
+  n02087394: Rhodesian_ridgeback
+  n02088094: Afghan_hound
+  n02088238: basset
+  n02088364: beagle
+  n02088466: bloodhound
+  n02088632: bluetick
+  n02089078: black-and-tan_coonhound
+  n02089867: Walker_hound
+  n02089973: English_foxhound
+  n02090379: redbone
+  n02090622: borzoi
+  n02090721: Irish_wolfhound
+  n02091032: Italian_greyhound
+  n02091134: whippet
+  n02091244: Ibizan_hound
+  n02091467: Norwegian_elkhound
+  n02091635: otterhound
+  n02091831: Saluki
+  n02092002: Scottish_deerhound
+  n02092339: Weimaraner
+  n02093256: Staffordshire_bullterrier
+  n02093428: American_Staffordshire_terrier
+  n02093647: Bedlington_terrier
+  n02093754: Border_terrier
+  n02093859: Kerry_blue_terrier
+  n02093991: Irish_terrier
+  n02094114: Norfolk_terrier
+  n02094258: Norwich_terrier
+  n02094433: Yorkshire_terrier
+  n02095314: wire-haired_fox_terrier
+  n02095570: Lakeland_terrier
+  n02095889: Sealyham_terrier
+  n02096051: Airedale
+  n02096177: cairn
+  n02096294: Australian_terrier
+  n02096437: Dandie_Dinmont
+  n02096585: Boston_bull
+  n02097047: miniature_schnauzer
+  n02097130: giant_schnauzer
+  n02097209: standard_schnauzer
+  n02097298: Scotch_terrier
+  n02097474: Tibetan_terrier
+  n02097658: silky_terrier
+  n02098105: soft-coated_wheaten_terrier
+  n02098286: West_Highland_white_terrier
+  n02098413: Lhasa
+  n02099267: flat-coated_retriever
+  n02099429: curly-coated_retriever
+  n02099601: golden_retriever
+  n02099712: Labrador_retriever
+  n02099849: Chesapeake_Bay_retriever
+  n02100236: German_short-haired_pointer
+  n02100583: vizsla
+  n02100735: English_setter
+  n02100877: Irish_setter
+  n02101006: Gordon_setter
+  n02101388: Brittany_spaniel
+  n02101556: clumber
+  n02102040: English_springer
+  n02102177: Welsh_springer_spaniel
+  n02102318: cocker_spaniel
+  n02102480: Sussex_spaniel
+  n02102973: Irish_water_spaniel
+  n02104029: kuvasz
+  n02104365: schipperke
+  n02105056: groenendael
+  n02105162: malinois
+  n02105251: briard
+  n02105412: kelpie
+  n02105505: komondor
+  n02105641: Old_English_sheepdog
+  n02105855: Shetland_sheepdog
+  n02106030: collie
+  n02106166: Border_collie
+  n02106382: Bouvier_des_Flandres
+  n02106550: Rottweiler
+  n02106662: German_shepherd
+  n02107142: Doberman
+  n02107312: miniature_pinscher
+  n02107574: Greater_Swiss_Mountain_dog
+  n02107683: Bernese_mountain_dog
+  n02107908: Appenzeller
+  n02108000: EntleBucher
+  n02108089: boxer
+  n02108422: bull_mastiff
+  n02108551: Tibetan_mastiff
+  n02108915: French_bulldog
+  n02109047: Great_Dane
+  n02109525: Saint_Bernard
+  n02109961: Eskimo_dog
+  n02110063: malamute
+  n02110185: Siberian_husky
+  n02110341: dalmatian
+  n02110627: affenpinscher
+  n02110806: basenji
+  n02110958: pug
+  n02111129: Leonberg
+  n02111277: Newfoundland
+  n02111500: Great_Pyrenees
+  n02111889: Samoyed
+  n02112018: Pomeranian
+  n02112137: chow
+  n02112350: keeshond
+  n02112706: Brabancon_griffon
+  n02113023: Pembroke
+  n02113186: Cardigan
+  n02113624: toy_poodle
+  n02113712: miniature_poodle
+  n02113799: standard_poodle
+  n02113978: Mexican_hairless
+  n02114367: timber_wolf
+  n02114548: white_wolf
+  n02114712: red_wolf
+  n02114855: coyote
+  n02115641: dingo
+  n02115913: dhole
+  n02116738: African_hunting_dog
+  n02117135: hyena
+  n02119022: red_fox
+  n02119789: kit_fox
+  n02120079: Arctic_fox
+  n02120505: grey_fox
+  n02123045: tabby
+  n02123159: tiger_cat
+  n02123394: Persian_cat
+  n02123597: Siamese_cat
+  n02124075: Egyptian_cat
+  n02125311: cougar
+  n02127052: lynx
+  n02128385: leopard
+  n02128757: snow_leopard
+  n02128925: jaguar
+  n02129165: lion
+  n02129604: tiger
+  n02130308: cheetah
+  n02132136: brown_bear
+  n02133161: American_black_bear
+  n02134084: ice_bear
+  n02134418: sloth_bear
+  n02137549: mongoose
+  n02138441: meerkat
+  n02165105: tiger_beetle
+  n02165456: ladybug
+  n02167151: ground_beetle
+  n02168699: long-horned_beetle
+  n02169497: leaf_beetle
+  n02172182: dung_beetle
+  n02174001: rhinoceros_beetle
+  n02177972: weevil
+  n02190166: fly
+  n02206856: bee
+  n02219486: ant
+  n02226429: grasshopper
+  n02229544: cricket
+  n02231487: walking_stick
+  n02233338: cockroach
+  n02236044: mantis
+  n02256656: cicada
+  n02259212: leafhopper
+  n02264363: lacewing
+  n02268443: dragonfly
+  n02268853: damselfly
+  n02276258: admiral
+  n02277742: ringlet
+  n02279972: monarch
+  n02280649: cabbage_butterfly
+  n02281406: sulphur_butterfly
+  n02281787: lycaenid
+  n02317335: starfish
+  n02319095: sea_urchin
+  n02321529: sea_cucumber
+  n02325366: wood_rabbit
+  n02326432: hare
+  n02328150: Angora
+  n02342885: hamster
+  n02346627: porcupine
+  n02356798: fox_squirrel
+  n02361337: marmot
+  n02363005: beaver
+  n02364673: guinea_pig
+  n02389026: sorrel
+  n02391049: zebra
+  n02395406: hog
+  n02396427: wild_boar
+  n02397096: warthog
+  n02398521: hippopotamus
+  n02403003: ox
+  n02408429: water_buffalo
+  n02410509: bison
+  n02412080: ram
+  n02415577: bighorn
+  n02417914: ibex
+  n02422106: hartebeest
+  n02422699: impala
+  n02423022: gazelle
+  n02437312: Arabian_camel
+  n02437616: llama
+  n02441942: weasel
+  n02442845: mink
+  n02443114: polecat
+  n02443484: black-footed_ferret
+  n02444819: otter
+  n02445715: skunk
+  n02447366: badger
+  n02454379: armadillo
+  n02457408: three-toed_sloth
+  n02480495: orangutan
+  n02480855: gorilla
+  n02481823: chimpanzee
+  n02483362: gibbon
+  n02483708: siamang
+  n02484975: guenon
+  n02486261: patas
+  n02486410: baboon
+  n02487347: macaque
+  n02488291: langur
+  n02488702: colobus
+  n02489166: proboscis_monkey
+  n02490219: marmoset
+  n02492035: capuchin
+  n02492660: howler_monkey
+  n02493509: titi
+  n02493793: spider_monkey
+  n02494079: squirrel_monkey
+  n02497673: Madagascar_cat
+  n02500267: indri
+  n02504013: Indian_elephant
+  n02504458: African_elephant
+  n02509815: lesser_panda
+  n02510455: giant_panda
+  n02514041: barracouta
+  n02526121: eel
+  n02536864: coho
+  n02606052: rock_beauty
+  n02607072: anemone_fish
+  n02640242: sturgeon
+  n02641379: gar
+  n02643566: lionfish
+  n02655020: puffer
+  n02666196: abacus
+  n02667093: abaya
+  n02669723: academic_gown
+  n02672831: accordion
+  n02676566: acoustic_guitar
+  n02687172: aircraft_carrier
+  n02690373: airliner
+  n02692877: airship
+  n02699494: altar
+  n02701002: ambulance
+  n02704792: amphibian
+  n02708093: analog_clock
+  n02727426: apiary
+  n02730930: apron
+  n02747177: ashcan
+  n02749479: assault_rifle
+  n02769748: backpack
+  n02776631: bakery
+  n02777292: balance_beam
+  n02782093: balloon
+  n02783161: ballpoint
+  n02786058: Band_Aid
+  n02787622: banjo
+  n02788148: bannister
+  n02790996: barbell
+  n02791124: barber_chair
+  n02791270: barbershop
+  n02793495: barn
+  n02794156: barometer
+  n02795169: barrel
+  n02797295: barrow
+  n02799071: baseball
+  n02802426: basketball
+  n02804414: bassinet
+  n02804610: bassoon
+  n02807133: bathing_cap
+  n02808304: bath_towel
+  n02808440: bathtub
+  n02814533: beach_wagon
+  n02814860: beacon
+  n02815834: beaker
+  n02817516: bearskin
+  n02823428: beer_bottle
+  n02823750: beer_glass
+  n02825657: bell_cote
+  n02834397: bib
+  n02835271: bicycle-built-for-two
+  n02837789: bikini
+  n02840245: binder
+  n02841315: binoculars
+  n02843684: birdhouse
+  n02859443: boathouse
+  n02860847: bobsled
+  n02865351: bolo_tie
+  n02869837: bonnet
+  n02870880: bookcase
+  n02871525: bookshop
+  n02877765: bottlecap
+  n02879718: bow
+  n02883205: bow_tie
+  n02892201: brass
+  n02892767: brassiere
+  n02894605: breakwater
+  n02895154: breastplate
+  n02906734: broom
+  n02909870: bucket
+  n02910353: buckle
+  n02916936: bulletproof_vest
+  n02917067: bullet_train
+  n02927161: butcher_shop
+  n02930766: cab
+  n02939185: caldron
+  n02948072: candle
+  n02950826: cannon
+  n02951358: canoe
+  n02951585: can_opener
+  n02963159: cardigan
+  n02965783: car_mirror
+  n02966193: carousel
+  n02966687: carpenter's_kit
+  n02971356: carton
+  n02974003: car_wheel
+  n02977058: cash_machine
+  n02978881: cassette
+  n02979186: cassette_player
+  n02980441: castle
+  n02981792: catamaran
+  n02988304: CD_player
+  n02992211: cello
+  n02992529: cellular_telephone
+  n02999410: chain
+  n03000134: chainlink_fence
+  n03000247: chain_mail
+  n03000684: chain_saw
+  n03014705: chest
+  n03016953: chiffonier
+  n03017168: chime
+  n03018349: china_cabinet
+  n03026506: Christmas_stocking
+  n03028079: church
+  n03032252: cinema
+  n03041632: cleaver
+  n03042490: cliff_dwelling
+  n03045698: cloak
+  n03047690: clog
+  n03062245: cocktail_shaker
+  n03063599: coffee_mug
+  n03063689: coffeepot
+  n03065424: coil
+  n03075370: combination_lock
+  n03085013: computer_keyboard
+  n03089624: confectionery
+  n03095699: container_ship
+  n03100240: convertible
+  n03109150: corkscrew
+  n03110669: cornet
+  n03124043: cowboy_boot
+  n03124170: cowboy_hat
+  n03125729: cradle
+  n03126707: crane_(machine)
+  n03127747: crash_helmet
+  n03127925: crate
+  n03131574: crib
+  n03133878: Crock_Pot
+  n03134739: croquet_ball
+  n03141823: crutch
+  n03146219: cuirass
+  n03160309: dam
+  n03179701: desk
+  n03180011: desktop_computer
+  n03187595: dial_telephone
+  n03188531: diaper
+  n03196217: digital_clock
+  n03197337: digital_watch
+  n03201208: dining_table
+  n03207743: dishrag
+  n03207941: dishwasher
+  n03208938: disk_brake
+  n03216828: dock
+  n03218198: dogsled
+  n03220513: dome
+  n03223299: doormat
+  n03240683: drilling_platform
+  n03249569: drum
+  n03250847: drumstick
+  n03255030: dumbbell
+  n03259280: Dutch_oven
+  n03271574: electric_fan
+  n03272010: electric_guitar
+  n03272562: electric_locomotive
+  n03290653: entertainment_center
+  n03291819: envelope
+  n03297495: espresso_maker
+  n03314780: face_powder
+  n03325584: feather_boa
+  n03337140: file
+  n03344393: fireboat
+  n03345487: fire_engine
+  n03347037: fire_screen
+  n03355925: flagpole
+  n03372029: flute
+  n03376595: folding_chair
+  n03379051: football_helmet
+  n03384352: forklift
+  n03388043: fountain
+  n03388183: fountain_pen
+  n03388549: four-poster
+  n03393912: freight_car
+  n03394916: French_horn
+  n03400231: frying_pan
+  n03404251: fur_coat
+  n03417042: garbage_truck
+  n03424325: gasmask
+  n03425413: gas_pump
+  n03443371: goblet
+  n03444034: go-kart
+  n03445777: golf_ball
+  n03445924: golfcart
+  n03447447: gondola
+  n03447721: gong
+  n03450230: gown
+  n03452741: grand_piano
+  n03457902: greenhouse
+  n03459775: grille
+  n03461385: grocery_store
+  n03467068: guillotine
+  n03476684: hair_slide
+  n03476991: hair_spray
+  n03478589: half_track
+  n03481172: hammer
+  n03482405: hamper
+  n03483316: hand_blower
+  n03485407: hand-held_computer
+  n03485794: handkerchief
+  n03492542: hard_disc
+  n03494278: harmonica
+  n03495258: harp
+  n03496892: harvester
+  n03498962: hatchet
+  n03527444: holster
+  n03529860: home_theater
+  n03530642: honeycomb
+  n03532672: hook
+  n03534580: hoopskirt
+  n03535780: horizontal_bar
+  n03538406: horse_cart
+  n03544143: hourglass
+  n03584254: iPod
+  n03584829: iron
+  n03590841: jack-o'-lantern
+  n03594734: jean
+  n03594945: jeep
+  n03595614: jersey
+  n03598930: jigsaw_puzzle
+  n03599486: jinrikisha
+  n03602883: joystick
+  n03617480: kimono
+  n03623198: knee_pad
+  n03627232: knot
+  n03630383: lab_coat
+  n03633091: ladle
+  n03637318: lampshade
+  n03642806: laptop
+  n03649909: lawn_mower
+  n03657121: lens_cap
+  n03658185: letter_opener
+  n03661043: library
+  n03662601: lifeboat
+  n03666591: lighter
+  n03670208: limousine
+  n03673027: liner
+  n03676483: lipstick
+  n03680355: Loafer
+  n03690938: lotion
+  n03691459: loudspeaker
+  n03692522: loupe
+  n03697007: lumbermill
+  n03706229: magnetic_compass
+  n03709823: mailbag
+  n03710193: mailbox
+  n03710637: maillot_(tights)
+  n03710721: maillot_(tank_suit)
+  n03717622: manhole_cover
+  n03720891: maraca
+  n03721384: marimba
+  n03724870: mask
+  n03729826: matchstick
+  n03733131: maypole
+  n03733281: maze
+  n03733805: measuring_cup
+  n03742115: medicine_chest
+  n03743016: megalith
+  n03759954: microphone
+  n03761084: microwave
+  n03763968: military_uniform
+  n03764736: milk_can
+  n03769881: minibus
+  n03770439: miniskirt
+  n03770679: minivan
+  n03773504: missile
+  n03775071: mitten
+  n03775546: mixing_bowl
+  n03776460: mobile_home
+  n03777568: Model_T
+  n03777754: modem
+  n03781244: monastery
+  n03782006: monitor
+  n03785016: moped
+  n03786901: mortar
+  n03787032: mortarboard
+  n03788195: mosque
+  n03788365: mosquito_net
+  n03791053: motor_scooter
+  n03792782: mountain_bike
+  n03792972: mountain_tent
+  n03793489: mouse
+  n03794056: mousetrap
+  n03796401: moving_van
+  n03803284: muzzle
+  n03804744: nail
+  n03814639: neck_brace
+  n03814906: necklace
+  n03825788: nipple
+  n03832673: notebook
+  n03837869: obelisk
+  n03838899: oboe
+  n03840681: ocarina
+  n03841143: odometer
+  n03843555: oil_filter
+  n03854065: organ
+  n03857828: oscilloscope
+  n03866082: overskirt
+  n03868242: oxcart
+  n03868863: oxygen_mask
+  n03871628: packet
+  n03873416: paddle
+  n03874293: paddlewheel
+  n03874599: padlock
+  n03876231: paintbrush
+  n03877472: pajama
+  n03877845: palace
+  n03884397: panpipe
+  n03887697: paper_towel
+  n03888257: parachute
+  n03888605: parallel_bars
+  n03891251: park_bench
+  n03891332: parking_meter
+  n03895866: passenger_car
+  n03899768: patio
+  n03902125: pay-phone
+  n03903868: pedestal
+  n03908618: pencil_box
+  n03908714: pencil_sharpener
+  n03916031: perfume
+  n03920288: Petri_dish
+  n03924679: photocopier
+  n03929660: pick
+  n03929855: pickelhaube
+  n03930313: picket_fence
+  n03930630: pickup
+  n03933933: pier
+  n03935335: piggy_bank
+  n03937543: pill_bottle
+  n03938244: pillow
+  n03942813: ping-pong_ball
+  n03944341: pinwheel
+  n03947888: pirate
+  n03950228: pitcher
+  n03954731: plane
+  n03956157: planetarium
+  n03958227: plastic_bag
+  n03961711: plate_rack
+  n03967562: plow
+  n03970156: plunger
+  n03976467: Polaroid_camera
+  n03976657: pole
+  n03977966: police_van
+  n03980874: poncho
+  n03982430: pool_table
+  n03983396: pop_bottle
+  n03991062: pot
+  n03992509: potter's_wheel
+  n03995372: power_drill
+  n03998194: prayer_rug
+  n04004767: printer
+  n04005630: prison
+  n04008634: projectile
+  n04009552: projector
+  n04019541: puck
+  n04023962: punching_bag
+  n04026417: purse
+  n04033901: quill
+  n04033995: quilt
+  n04037443: racer
+  n04039381: racket
+  n04040759: radiator
+  n04041544: radio
+  n04044716: radio_telescope
+  n04049303: rain_barrel
+  n04065272: recreational_vehicle
+  n04067472: reel
+  n04069434: reflex_camera
+  n04070727: refrigerator
+  n04074963: remote_control
+  n04081281: restaurant
+  n04086273: revolver
+  n04090263: rifle
+  n04099969: rocking_chair
+  n04111531: rotisserie
+  n04116512: rubber_eraser
+  n04118538: rugby_ball
+  n04118776: rule
+  n04120489: running_shoe
+  n04125021: safe
+  n04127249: safety_pin
+  n04131690: saltshaker
+  n04133789: sandal
+  n04136333: sarong
+  n04141076: sax
+  n04141327: scabbard
+  n04141975: scale
+  n04146614: school_bus
+  n04147183: schooner
+  n04149813: scoreboard
+  n04152593: screen
+  n04153751: screw
+  n04154565: screwdriver
+  n04162706: seat_belt
+  n04179913: sewing_machine
+  n04192698: shield
+  n04200800: shoe_shop
+  n04201297: shoji
+  n04204238: shopping_basket
+  n04204347: shopping_cart
+  n04208210: shovel
+  n04209133: shower_cap
+  n04209239: shower_curtain
+  n04228054: ski
+  n04229816: ski_mask
+  n04235860: sleeping_bag
+  n04238763: slide_rule
+  n04239074: sliding_door
+  n04243546: slot
+  n04251144: snorkel
+  n04252077: snowmobile
+  n04252225: snowplow
+  n04254120: soap_dispenser
+  n04254680: soccer_ball
+  n04254777: sock
+  n04258138: solar_dish
+  n04259630: sombrero
+  n04263257: soup_bowl
+  n04264628: space_bar
+  n04265275: space_heater
+  n04266014: space_shuttle
+  n04270147: spatula
+  n04273569: speedboat
+  n04275548: spider_web
+  n04277352: spindle
+  n04285008: sports_car
+  n04286575: spotlight
+  n04296562: stage
+  n04310018: steam_locomotive
+  n04311004: steel_arch_bridge
+  n04311174: steel_drum
+  n04317175: stethoscope
+  n04325704: stole
+  n04326547: stone_wall
+  n04328186: stopwatch
+  n04330267: stove
+  n04332243: strainer
+  n04335435: streetcar
+  n04336792: stretcher
+  n04344873: studio_couch
+  n04346328: stupa
+  n04347754: submarine
+  n04350905: suit
+  n04355338: sundial
+  n04355933: sunglass
+  n04356056: sunglasses
+  n04357314: sunscreen
+  n04366367: suspension_bridge
+  n04367480: swab
+  n04370456: sweatshirt
+  n04371430: swimming_trunks
+  n04371774: swing
+  n04372370: switch
+  n04376876: syringe
+  n04380533: table_lamp
+  n04389033: tank
+  n04392985: tape_player
+  n04398044: teapot
+  n04399382: teddy
+  n04404412: television
+  n04409515: tennis_ball
+  n04417672: thatch
+  n04418357: theater_curtain
+  n04423845: thimble
+  n04428191: thresher
+  n04429376: throne
+  n04435653: tile_roof
+  n04442312: toaster
+  n04443257: tobacco_shop
+  n04447861: toilet_seat
+  n04456115: torch
+  n04458633: totem_pole
+  n04461696: tow_truck
+  n04462240: toyshop
+  n04465501: tractor
+  n04467665: trailer_truck
+  n04476259: tray
+  n04479046: trench_coat
+  n04482393: tricycle
+  n04483307: trimaran
+  n04485082: tripod
+  n04486054: triumphal_arch
+  n04487081: trolleybus
+  n04487394: trombone
+  n04493381: tub
+  n04501370: turnstile
+  n04505470: typewriter_keyboard
+  n04507155: umbrella
+  n04509417: unicycle
+  n04515003: upright
+  n04517823: vacuum
+  n04522168: vase
+  n04523525: vault
+  n04525038: velvet
+  n04525305: vending_machine
+  n04532106: vestment
+  n04532670: viaduct
+  n04536866: violin
+  n04540053: volleyball
+  n04542943: waffle_iron
+  n04548280: wall_clock
+  n04548362: wallet
+  n04550184: wardrobe
+  n04552348: warplane
+  n04553703: washbasin
+  n04554684: washer
+  n04557648: water_bottle
+  n04560804: water_jug
+  n04562935: water_tower
+  n04579145: whiskey_jug
+  n04579432: whistle
+  n04584207: wig
+  n04589890: window_screen
+  n04590129: window_shade
+  n04591157: Windsor_tie
+  n04591713: wine_bottle
+  n04592741: wing
+  n04596742: wok
+  n04597913: wooden_spoon
+  n04599235: wool
+  n04604644: worm_fence
+  n04606251: wreck
+  n04612504: yawl
+  n04613696: yurt
+  n06359193: web_site
+  n06596364: comic_book
+  n06785654: crossword_puzzle
+  n06794110: street_sign
+  n06874185: traffic_light
+  n07248320: book_jacket
+  n07565083: menu
+  n07579787: plate
+  n07583066: guacamole
+  n07584110: consomme
+  n07590611: hot_pot
+  n07613480: trifle
+  n07614500: ice_cream
+  n07615774: ice_lolly
+  n07684084: French_loaf
+  n07693725: bagel
+  n07695742: pretzel
+  n07697313: cheeseburger
+  n07697537: hotdog
+  n07711569: mashed_potato
+  n07714571: head_cabbage
+  n07714990: broccoli
+  n07715103: cauliflower
+  n07716358: zucchini
+  n07716906: spaghetti_squash
+  n07717410: acorn_squash
+  n07717556: butternut_squash
+  n07718472: cucumber
+  n07718747: artichoke
+  n07720875: bell_pepper
+  n07730033: cardoon
+  n07734744: mushroom
+  n07742313: Granny_Smith
+  n07745940: strawberry
+  n07747607: orange
+  n07749582: lemon
+  n07753113: fig
+  n07753275: pineapple
+  n07753592: banana
+  n07754684: jackfruit
+  n07760859: custard_apple
+  n07768694: pomegranate
+  n07802026: hay
+  n07831146: carbonara
+  n07836838: chocolate_sauce
+  n07860988: dough
+  n07871810: meat_loaf
+  n07873807: pizza
+  n07875152: potpie
+  n07880968: burrito
+  n07892512: red_wine
+  n07920052: espresso
+  n07930864: cup
+  n07932039: eggnog
+  n09193705: alp
+  n09229709: bubble
+  n09246464: cliff
+  n09256479: coral_reef
+  n09288635: geyser
+  n09332890: lakeside
+  n09399592: promontory
+  n09421951: sandbar
+  n09428293: seashore
+  n09468604: valley
+  n09472597: volcano
+  n09835506: ballplayer
+  n10148035: groom
+  n10565667: scuba_diver
+  n11879895: rapeseed
+  n11939491: daisy
+  n12057211: yellow_lady's_slipper
+  n12144580: corn
+  n12267677: acorn
+  n12620546: hip
+  n12768682: buckeye
+  n12985857: coral_fungus
+  n12998815: agaric
+  n13037406: gyromitra
+  n13040303: stinkhorn
+  n13044778: earthstar
+  n13052670: hen-of-the-woods
+  n13054560: bolete
+  n13133613: ear
+  n15075141: toilet_tissue
+
+
+# Download script/URL (optional)
+download: yolo/data/scripts/get_imagenet.sh

ultralytics/datasets/Objects365.yaml

@@ -0,0 +1,443 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Objects365 dataset https://www.objects365.org/ by Megvii
+# Example usage: yolo train data=Objects365.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── Objects365  ← downloads here (712 GB = 367G data + 345G zips)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/Objects365  # dataset root dir
+train: images/train  # train images (relative to 'path') 1742289 images
+val: images/val # val images (relative to 'path') 80000 images
+test:  # test images (optional)
+
+# Classes
+names:
+  0: Person
+  1: Sneakers
+  2: Chair
+  3: Other Shoes
+  4: Hat
+  5: Car
+  6: Lamp
+  7: Glasses
+  8: Bottle
+  9: Desk
+  10: Cup
+  11: Street Lights
+  12: Cabinet/shelf
+  13: Handbag/Satchel
+  14: Bracelet
+  15: Plate
+  16: Picture/Frame
+  17: Helmet
+  18: Book
+  19: Gloves
+  20: Storage box
+  21: Boat
+  22: Leather Shoes
+  23: Flower
+  24: Bench
+  25: Potted Plant
+  26: Bowl/Basin
+  27: Flag
+  28: Pillow
+  29: Boots
+  30: Vase
+  31: Microphone
+  32: Necklace
+  33: Ring
+  34: SUV
+  35: Wine Glass
+  36: Belt
+  37: Monitor/TV
+  38: Backpack
+  39: Umbrella
+  40: Traffic Light
+  41: Speaker
+  42: Watch
+  43: Tie
+  44: Trash bin Can
+  45: Slippers
+  46: Bicycle
+  47: Stool
+  48: Barrel/bucket
+  49: Van
+  50: Couch
+  51: Sandals
+  52: Basket
+  53: Drum
+  54: Pen/Pencil
+  55: Bus
+  56: Wild Bird
+  57: High Heels
+  58: Motorcycle
+  59: Guitar
+  60: Carpet
+  61: Cell Phone
+  62: Bread
+  63: Camera
+  64: Canned
+  65: Truck
+  66: Traffic cone
+  67: Cymbal
+  68: Lifesaver
+  69: Towel
+  70: Stuffed Toy
+  71: Candle
+  72: Sailboat
+  73: Laptop
+  74: Awning
+  75: Bed
+  76: Faucet
+  77: Tent
+  78: Horse
+  79: Mirror
+  80: Power outlet
+  81: Sink
+  82: Apple
+  83: Air Conditioner
+  84: Knife
+  85: Hockey Stick
+  86: Paddle
+  87: Pickup Truck
+  88: Fork
+  89: Traffic Sign
+  90: Balloon
+  91: Tripod
+  92: Dog
+  93: Spoon
+  94: Clock
+  95: Pot
+  96: Cow
+  97: Cake
+  98: Dinning Table
+  99: Sheep
+  100: Hanger
+  101: Blackboard/Whiteboard
+  102: Napkin
+  103: Other Fish
+  104: Orange/Tangerine
+  105: Toiletry
+  106: Keyboard
+  107: Tomato
+  108: Lantern
+  109: Machinery Vehicle
+  110: Fan
+  111: Green Vegetables
+  112: Banana
+  113: Baseball Glove
+  114: Airplane
+  115: Mouse
+  116: Train
+  117: Pumpkin
+  118: Soccer
+  119: Skiboard
+  120: Luggage
+  121: Nightstand
+  122: Tea pot
+  123: Telephone
+  124: Trolley
+  125: Head Phone
+  126: Sports Car
+  127: Stop Sign
+  128: Dessert
+  129: Scooter
+  130: Stroller
+  131: Crane
+  132: Remote
+  133: Refrigerator
+  134: Oven
+  135: Lemon
+  136: Duck
+  137: Baseball Bat
+  138: Surveillance Camera
+  139: Cat
+  140: Jug
+  141: Broccoli
+  142: Piano
+  143: Pizza
+  144: Elephant
+  145: Skateboard
+  146: Surfboard
+  147: Gun
+  148: Skating and Skiing shoes
+  149: Gas stove
+  150: Donut
+  151: Bow Tie
+  152: Carrot
+  153: Toilet
+  154: Kite
+  155: Strawberry
+  156: Other Balls
+  157: Shovel
+  158: Pepper
+  159: Computer Box
+  160: Toilet Paper
+  161: Cleaning Products
+  162: Chopsticks
+  163: Microwave
+  164: Pigeon
+  165: Baseball
+  166: Cutting/chopping Board
+  167: Coffee Table
+  168: Side Table
+  169: Scissors
+  170: Marker
+  171: Pie
+  172: Ladder
+  173: Snowboard
+  174: Cookies
+  175: Radiator
+  176: Fire Hydrant
+  177: Basketball
+  178: Zebra
+  179: Grape
+  180: Giraffe
+  181: Potato
+  182: Sausage
+  183: Tricycle
+  184: Violin
+  185: Egg
+  186: Fire Extinguisher
+  187: Candy
+  188: Fire Truck
+  189: Billiards
+  190: Converter
+  191: Bathtub
+  192: Wheelchair
+  193: Golf Club
+  194: Briefcase
+  195: Cucumber
+  196: Cigar/Cigarette
+  197: Paint Brush
+  198: Pear
+  199: Heavy Truck
+  200: Hamburger
+  201: Extractor
+  202: Extension Cord
+  203: Tong
+  204: Tennis Racket
+  205: Folder
+  206: American Football
+  207: earphone
+  208: Mask
+  209: Kettle
+  210: Tennis
+  211: Ship
+  212: Swing
+  213: Coffee Machine
+  214: Slide
+  215: Carriage
+  216: Onion
+  217: Green beans
+  218: Projector
+  219: Frisbee
+  220: Washing Machine/Drying Machine
+  221: Chicken
+  222: Printer
+  223: Watermelon
+  224: Saxophone
+  225: Tissue
+  226: Toothbrush
+  227: Ice cream
+  228: Hot-air balloon
+  229: Cello
+  230: French Fries
+  231: Scale
+  232: Trophy
+  233: Cabbage
+  234: Hot dog
+  235: Blender
+  236: Peach
+  237: Rice
+  238: Wallet/Purse
+  239: Volleyball
+  240: Deer
+  241: Goose
+  242: Tape
+  243: Tablet
+  244: Cosmetics
+  245: Trumpet
+  246: Pineapple
+  247: Golf Ball
+  248: Ambulance
+  249: Parking meter
+  250: Mango
+  251: Key
+  252: Hurdle
+  253: Fishing Rod
+  254: Medal
+  255: Flute
+  256: Brush
+  257: Penguin
+  258: Megaphone
+  259: Corn
+  260: Lettuce
+  261: Garlic
+  262: Swan
+  263: Helicopter
+  264: Green Onion
+  265: Sandwich
+  266: Nuts
+  267: Speed Limit Sign
+  268: Induction Cooker
+  269: Broom
+  270: Trombone
+  271: Plum
+  272: Rickshaw
+  273: Goldfish
+  274: Kiwi fruit
+  275: Router/modem
+  276: Poker Card
+  277: Toaster
+  278: Shrimp
+  279: Sushi
+  280: Cheese
+  281: Notepaper
+  282: Cherry
+  283: Pliers
+  284: CD
+  285: Pasta
+  286: Hammer
+  287: Cue
+  288: Avocado
+  289: Hamimelon
+  290: Flask
+  291: Mushroom
+  292: Screwdriver
+  293: Soap
+  294: Recorder
+  295: Bear
+  296: Eggplant
+  297: Board Eraser
+  298: Coconut
+  299: Tape Measure/Ruler
+  300: Pig
+  301: Showerhead
+  302: Globe
+  303: Chips
+  304: Steak
+  305: Crosswalk Sign
+  306: Stapler
+  307: Camel
+  308: Formula 1
+  309: Pomegranate
+  310: Dishwasher
+  311: Crab
+  312: Hoverboard
+  313: Meat ball
+  314: Rice Cooker
+  315: Tuba
+  316: Calculator
+  317: Papaya
+  318: Antelope
+  319: Parrot
+  320: Seal
+  321: Butterfly
+  322: Dumbbell
+  323: Donkey
+  324: Lion
+  325: Urinal
+  326: Dolphin
+  327: Electric Drill
+  328: Hair Dryer
+  329: Egg tart
+  330: Jellyfish
+  331: Treadmill
+  332: Lighter
+  333: Grapefruit
+  334: Game board
+  335: Mop
+  336: Radish
+  337: Baozi
+  338: Target
+  339: French
+  340: Spring Rolls
+  341: Monkey
+  342: Rabbit
+  343: Pencil Case
+  344: Yak
+  345: Red Cabbage
+  346: Binoculars
+  347: Asparagus
+  348: Barbell
+  349: Scallop
+  350: Noddles
+  351: Comb
+  352: Dumpling
+  353: Oyster
+  354: Table Tennis paddle
+  355: Cosmetics Brush/Eyeliner Pencil
+  356: Chainsaw
+  357: Eraser
+  358: Lobster
+  359: Durian
+  360: Okra
+  361: Lipstick
+  362: Cosmetics Mirror
+  363: Curling
+  364: Table Tennis
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  from tqdm import tqdm
+
+  from ultralytics.yolo.utils.checks import check_requirements
+  from ultralytics.yolo.utils.downloads import download
+  from ultralytics.yolo.utils.ops import xyxy2xywhn
+
+  import numpy as np
+  from pathlib import Path
+
+  check_requirements(('pycocotools>=2.0',))
+  from pycocotools.coco import COCO
+
+  # Make Directories
+  dir = Path(yaml['path'])  # dataset root dir
+  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)
+
+  # Train, Val Splits
+  for split, patches in [('train', 50 + 1), ('val', 43 + 1)]:
+      print(f"Processing {split} in {patches} patches ...")
+      images, labels = dir / 'images' / split, dir / 'labels' / split
+
+      # Download
+      url = f"https://dorc.ks3-cn-beijing.ksyun.com/data-set/2020Objects365%E6%95%B0%E6%8D%AE%E9%9B%86/{split}/"
+      if split == 'train':
+          download([f'{url}zhiyuan_objv2_{split}.tar.gz'], dir=dir)  # annotations json
+          download([f'{url}patch{i}.tar.gz' for i in range(patches)], dir=images, curl=True, threads=8)
+      elif split == 'val':
+          download([f'{url}zhiyuan_objv2_{split}.json'], dir=dir)  # annotations json
+          download([f'{url}images/v1/patch{i}.tar.gz' for i in range(15 + 1)], dir=images, curl=True, threads=8)
+          download([f'{url}images/v2/patch{i}.tar.gz' for i in range(16, patches)], dir=images, curl=True, threads=8)
+
+      # Move
+      for f in tqdm(images.rglob('*.jpg'), desc=f'Moving {split} images'):
+          f.rename(images / f.name)  # move to /images/{split}
+
+      # Labels
+      coco = COCO(dir / f'zhiyuan_objv2_{split}.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(labels / 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)
+                          xyxy = np.array([x, y, x + w, y + h])[None]  # pixels(1,4)
+                          x, y, w, h = xyxy2xywhn(xyxy, w=width, h=height, clip=True)[0]  # normalized and clipped
+                          file.write(f"{cid} {x:.5f} {y:.5f} {w:.5f} {h:.5f}\n")
+              except Exception as e:
+                  print(e)

ultralytics/datasets/SKU-110K.yaml

@@ -0,0 +1,58 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19 by Trax Retail
+# Example usage: yolo train data=SKU-110K.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── SKU-110K  ← downloads here (13.6 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/SKU-110K  # dataset root dir
+train: train.txt  # train images (relative to 'path')  8219 images
+val: val.txt  # val images (relative to 'path')  588 images
+test: test.txt  # test images (optional)  2936 images
+
+# Classes
+names:
+  0: object
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import shutil
+  from pathlib import Path
+
+  import numpy as np
+  import pandas as pd
+  from tqdm import tqdm
+
+  from ultralytics.yolo.utils.downloads import download
+  from ultralytics.yolo.utils.ops import xyxy2xywh
+
+  # Download
+  dir = Path(yaml['path'])  # dataset root dir
+  parent = Path(dir.parent)  # download dir
+  urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz']
+  download(urls, dir=parent)
+
+  # Rename directories
+  if dir.exists():
+      shutil.rmtree(dir)
+  (parent / '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

ultralytics/datasets/VOC.yaml

@@ -0,0 +1,100 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
+# Example usage: yolo train data=VOC.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── VOC  ← downloads here (2.8 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/VOC
+train: # train images (relative to 'path')  16551 images
+  - images/train2012
+  - images/train2007
+  - images/val2012
+  - images/val2007
+val: # val images (relative to 'path')  4952 images
+  - images/test2007
+test: # test images (optional)
+  - images/test2007
+
+# Classes
+names:
+  0: aeroplane
+  1: bicycle
+  2: bird
+  3: boat
+  4: bottle
+  5: bus
+  6: car
+  7: cat
+  8: chair
+  9: cow
+  10: diningtable
+  11: dog
+  12: horse
+  13: motorbike
+  14: person
+  15: pottedplant
+  16: sheep
+  17: sofa
+  18: train
+  19: tvmonitor
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import xml.etree.ElementTree as ET
+
+  from tqdm import tqdm
+  from ultralytics.yolo.utils.downloads import download
+  from pathlib import Path
+
+  def convert_label(path, lb_path, year, image_id):
+      def convert_box(size, box):
+          dw, dh = 1. / size[0], 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
+
+      in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')
+      out_file = open(lb_path, 'w')
+      tree = ET.parse(in_file)
+      root = tree.getroot()
+      size = root.find('size')
+      w = int(size.find('width').text)
+      h = int(size.find('height').text)
+
+      names = list(yaml['names'].values())  # names list
+      for obj in root.iter('object'):
+          cls = obj.find('name').text
+          if cls in names and int(obj.find('difficult').text) != 1:
+              xmlbox = obj.find('bndbox')
+              bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
+              cls_id = names.index(cls)  # class id
+              out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')
+
+
+  # Download
+  dir = Path(yaml['path'])  # dataset root dir
+  url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
+  urls = [f'{url}VOCtrainval_06-Nov-2007.zip',  # 446MB, 5012 images
+          f'{url}VOCtest_06-Nov-2007.zip',  # 438MB, 4953 images
+          f'{url}VOCtrainval_11-May-2012.zip']  # 1.95GB, 17126 images
+  download(urls, dir=dir / 'images', curl=True, threads=3)
+
+  # Convert
+  path = dir / 'images/VOCdevkit'
+  for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'):
+      imgs_path = dir / 'images' / f'{image_set}{year}'
+      lbs_path = dir / 'labels' / f'{image_set}{year}'
+      imgs_path.mkdir(exist_ok=True, parents=True)
+      lbs_path.mkdir(exist_ok=True, parents=True)
+
+      with open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt') as f:
+          image_ids = f.read().strip().split()
+      for id in tqdm(image_ids, desc=f'{image_set}{year}'):
+          f = path / f'VOC{year}/JPEGImages/{id}.jpg'  # old img path
+          lb_path = (lbs_path / f.name).with_suffix('.txt')  # new label path
+          f.rename(imgs_path / f.name)  # move image
+          convert_label(path, lb_path, year, id)  # convert labels to YOLO format

ultralytics/datasets/VisDrone.yaml

@@ -0,0 +1,73 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University
+# Example usage: yolo train data=VisDrone.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── VisDrone  ← downloads here (2.3 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/VisDrone  # dataset root dir
+train: VisDrone2019-DET-train/images  # train images (relative to 'path')  6471 images
+val: VisDrone2019-DET-val/images  # val images (relative to 'path')  548 images
+test: VisDrone2019-DET-test-dev/images  # test images (optional)  1610 images
+
+# Classes
+names:
+  0: pedestrian
+  1: people
+  2: bicycle
+  3: car
+  4: van
+  5: truck
+  6: tricycle
+  7: awning-tricycle
+  8: bus
+  9: motor
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import os
+  from pathlib import Path
+
+  from ultralytics.yolo.utils.downloads import download
+
+  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(f'{os.sep}annotations{os.sep}', f'{os.sep}labels{os.sep}'), 'w') as fl:
+                      fl.writelines(lines)  # write label.txt
+
+
+  # Download
+  dir = Path(yaml['path'])  # dataset root dir
+  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, curl=True, threads=4)
+
+  # Convert
+  for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
+      visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels

ultralytics/datasets/coco-pose.yaml

@@ -0,0 +1,38 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO 2017 dataset http://cocodataset.org by Microsoft
+# Example usage: yolo train data=coco-pose.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco-pose  ← downloads here (20.1 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco-pose  # dataset root dir
+train: train2017.txt  # train images (relative to 'path') 118287 images
+val: val2017.txt  # val images (relative to 'path') 5000 images
+test: test-dev2017.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
+
+# Keypoints
+kpt_shape: [17, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
+flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+
+# Classes
+names:
+  0: person
+
+# Download script/URL (optional)
+download: |
+  from ultralytics.yolo.utils.downloads import download
+  from pathlib import Path
+
+  # Download labels
+  dir = Path(yaml['path'])  # dataset root dir
+  url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
+  urls = [url + 'coco2017labels-pose.zip']  # labels
+  download(urls, dir=dir.parent)
+  # Download data
+  urls = ['http://images.cocodataset.org/zips/train2017.zip',  # 19G, 118k images
+          'http://images.cocodataset.org/zips/val2017.zip',  # 1G, 5k images
+          'http://images.cocodataset.org/zips/test2017.zip']  # 7G, 41k images (optional)
+  download(urls, dir=dir / 'images', threads=3)

ultralytics/datasets/coco.yaml

@@ -0,0 +1,115 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO 2017 dataset http://cocodataset.org by Microsoft
+# Example usage: yolo train data=coco.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco  ← downloads here (20.1 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco  # dataset root dir
+train: train2017.txt  # train images (relative to 'path') 118287 images
+val: val2017.txt  # val images (relative to 'path') 5000 images
+test: test-dev2017.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush
+
+
+# Download script/URL (optional)
+download: |
+  from ultralytics.yolo.utils.downloads import download
+  from pathlib import Path
+
+  # Download labels
+  segments = True  # segment or box labels
+  dir = Path(yaml['path'])  # dataset root dir
+  url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
+  urls = [url + ('coco2017labels-segments.zip' if segments else 'coco2017labels.zip')]  # labels
+  download(urls, dir=dir.parent)
+  # Download data
+  urls = ['http://images.cocodataset.org/zips/train2017.zip',  # 19G, 118k images
+          'http://images.cocodataset.org/zips/val2017.zip',  # 1G, 5k images
+          'http://images.cocodataset.org/zips/test2017.zip']  # 7G, 41k images (optional)
+  download(urls, dir=dir / 'images', threads=3)

ultralytics/datasets/coco128-seg.yaml

@@ -0,0 +1,101 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
+# Example usage: yolo train data=coco128.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco128-seg  ← downloads here (7 MB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco128-seg  # dataset root dir
+train: images/train2017  # train images (relative to 'path') 128 images
+val: images/train2017  # val images (relative to 'path') 128 images
+test:  # test images (optional)
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush
+
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/coco128-seg.zip

ultralytics/datasets/coco128.yaml

@@ -0,0 +1,101 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
+# Example usage: yolo train data=coco128.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco128  ← downloads here (7 MB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco128  # dataset root dir
+train: images/train2017  # train images (relative to 'path') 128 images
+val: images/train2017  # val images (relative to 'path') 128 images
+test:  # test images (optional)
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush
+
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/coco128.zip

ultralytics/datasets/coco8-pose.yaml

@@ -0,0 +1,25 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO8-pose dataset (first 8 images from COCO train2017) by Ultralytics
+# Example usage: yolo train data=coco8-pose.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco8-pose  ← downloads here (1 MB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco8-pose  # dataset root dir
+train: images/train  # train images (relative to 'path') 4 images
+val: images/val  # val images (relative to 'path') 4 images
+test:  # test images (optional)
+
+# Keypoints
+kpt_shape: [17, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
+flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+
+# Classes
+names:
+  0: person
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/coco8-pose.zip

ultralytics/datasets/coco8-seg.yaml

@@ -0,0 +1,101 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO8-seg dataset (first 8 images from COCO train2017) by Ultralytics
+# Example usage: yolo train data=coco8-seg.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco8-seg  ← downloads here (1 MB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco8-seg  # dataset root dir
+train: images/train  # train images (relative to 'path') 4 images
+val: images/val  # val images (relative to 'path') 4 images
+test:  # test images (optional)
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush
+
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/coco8-seg.zip

ultralytics/datasets/coco8.yaml

@@ -0,0 +1,101 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# COCO8 dataset (first 8 images from COCO train2017) by Ultralytics
+# Example usage: yolo train data=coco8.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco8  ← downloads here (1 MB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/coco8  # dataset root dir
+train: images/train  # train images (relative to 'path') 4 images
+val: images/val  # val images (relative to 'path') 4 images
+test:  # test images (optional)
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush
+
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/coco8.zip

ultralytics/datasets/xView.yaml

@@ -0,0 +1,153 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# DIUx xView 2018 Challenge https://challenge.xviewdataset.org by U.S. National Geospatial-Intelligence Agency (NGA)
+# --------  DOWNLOAD DATA MANUALLY and jar xf val_images.zip to 'datasets/xView' before running train command!  --------
+# Example usage: yolo train data=xView.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── xView  ← downloads here (20.7 GB)
+
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/xView  # dataset root dir
+train: images/autosplit_train.txt  # train images (relative to 'path') 90% of 847 train images
+val: images/autosplit_val.txt  # train images (relative to 'path') 10% of 847 train images
+
+# Classes
+names:
+  0: Fixed-wing Aircraft
+  1: Small Aircraft
+  2: Cargo Plane
+  3: Helicopter
+  4: Passenger Vehicle
+  5: Small Car
+  6: Bus
+  7: Pickup Truck
+  8: Utility Truck
+  9: Truck
+  10: Cargo Truck
+  11: Truck w/Box
+  12: Truck Tractor
+  13: Trailer
+  14: Truck w/Flatbed
+  15: Truck w/Liquid
+  16: Crane Truck
+  17: Railway Vehicle
+  18: Passenger Car
+  19: Cargo Car
+  20: Flat Car
+  21: Tank car
+  22: Locomotive
+  23: Maritime Vessel
+  24: Motorboat
+  25: Sailboat
+  26: Tugboat
+  27: Barge
+  28: Fishing Vessel
+  29: Ferry
+  30: Yacht
+  31: Container Ship
+  32: Oil Tanker
+  33: Engineering Vehicle
+  34: Tower crane
+  35: Container Crane
+  36: Reach Stacker
+  37: Straddle Carrier
+  38: Mobile Crane
+  39: Dump Truck
+  40: Haul Truck
+  41: Scraper/Tractor
+  42: Front loader/Bulldozer
+  43: Excavator
+  44: Cement Mixer
+  45: Ground Grader
+  46: Hut/Tent
+  47: Shed
+  48: Building
+  49: Aircraft Hangar
+  50: Damaged Building
+  51: Facility
+  52: Construction Site
+  53: Vehicle Lot
+  54: Helipad
+  55: Storage Tank
+  56: Shipping container lot
+  57: Shipping Container
+  58: Pylon
+  59: Tower
+
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import json
+  import os
+  from pathlib import Path
+
+  import numpy as np
+  from PIL import Image
+  from tqdm import tqdm
+
+  from ultralytics.yolo.data.dataloaders.v5loader import autosplit
+  from ultralytics.yolo.utils.ops import xyxy2xywhn
+
+
+  def convert_labels(fname=Path('xView/xView_train.geojson')):
+      # Convert xView geoJSON labels to YOLO format
+      path = fname.parent
+      with open(fname) as f:
+          print(f'Loading {fname}...')
+          data = json.load(f)
+
+      # Make dirs
+      labels = Path(path / 'labels' / 'train')
+      os.system(f'rm -rf {labels}')
+      labels.mkdir(parents=True, exist_ok=True)
+
+      # xView classes 11-94 to 0-59
+      xview_class2index = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 0, 1, 2, -1, 3, -1, 4, 5, 6, 7, 8, -1, 9, 10, 11,
+                           12, 13, 14, 15, -1, -1, 16, 17, 18, 19, 20, 21, 22, -1, 23, 24, 25, -1, 26, 27, -1, 28, -1,
+                           29, 30, 31, 32, 33, 34, 35, 36, 37, -1, 38, 39, 40, 41, 42, 43, 44, 45, -1, -1, -1, -1, 46,
+                           47, 48, 49, -1, 50, 51, -1, 52, -1, -1, -1, 53, 54, -1, 55, -1, -1, 56, -1, 57, -1, 58, 59]
+
+      shapes = {}
+      for feature in tqdm(data['features'], desc=f'Converting {fname}'):
+          p = feature['properties']
+          if p['bounds_imcoords']:
+              id = p['image_id']
+              file = path / 'train_images' / id
+              if file.exists():  # 1395.tif missing
+                  try:
+                      box = np.array([int(num) for num in p['bounds_imcoords'].split(",")])
+                      assert box.shape[0] == 4, f'incorrect box shape {box.shape[0]}'
+                      cls = p['type_id']
+                      cls = xview_class2index[int(cls)]  # xView class to 0-60
+                      assert 59 >= cls >= 0, f'incorrect class index {cls}'
+
+                      # Write YOLO label
+                      if id not in shapes:
+                          shapes[id] = Image.open(file).size
+                      box = xyxy2xywhn(box[None].astype(np.float), w=shapes[id][0], h=shapes[id][1], clip=True)
+                      with open((labels / id).with_suffix('.txt'), 'a') as f:
+                          f.write(f"{cls} {' '.join(f'{x:.6f}' for x in box[0])}\n")  # write label.txt
+                  except Exception as e:
+                      print(f'WARNING: skipping one label for {file}: {e}')
+
+
+  # Download manually from https://challenge.xviewdataset.org
+  dir = Path(yaml['path'])  # dataset root dir
+  # urls = ['https://d307kc0mrhucc3.cloudfront.net/train_labels.zip',  # train labels
+  #         'https://d307kc0mrhucc3.cloudfront.net/train_images.zip',  # 15G, 847 train images
+  #         'https://d307kc0mrhucc3.cloudfront.net/val_images.zip']  # 5G, 282 val images (no labels)
+  # download(urls, dir=dir)
+
+  # Convert labels
+  convert_labels(dir / 'xView_train.geojson')
+
+  # Move images
+  images = Path(dir / 'images')
+  images.mkdir(parents=True, exist_ok=True)
+  Path(dir / 'train_images').rename(dir / 'images' / 'train')
+  Path(dir / 'val_images').rename(dir / 'images' / 'val')
+
+  # Split
+  autosplit(dir / 'images' / 'train')

ultralytics/hub/__init__.py

@@ -0,0 +1,117 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import requests
+
+from ultralytics.hub.auth import Auth
+from ultralytics.hub.utils import PREFIX
+from ultralytics.yolo.data.utils import HUBDatasetStats
+from ultralytics.yolo.utils import LOGGER, SETTINGS, USER_CONFIG_DIR, yaml_save
+
+
+def login(api_key=''):
+    """
+    Log in to the Ultralytics HUB API using the provided API key.
+
+    Args:
+        api_key (str, optional): May be an API key or a combination API key and model ID, i.e. key_id
+
+    Example:
+        from ultralytics import hub
+        hub.login('API_KEY')
+    """
+    Auth(api_key, verbose=True)
+
+
+def logout():
+    """
+    Log out of Ultralytics HUB by removing the API key from the settings file. To log in again, use 'yolo hub login'.
+
+    Example:
+        from ultralytics import hub
+        hub.logout()
+    """
+    SETTINGS['api_key'] = ''
+    yaml_save(USER_CONFIG_DIR / 'settings.yaml', SETTINGS)
+    LOGGER.info(f"{PREFIX}logged out ✅. To log in again, use 'yolo hub login'.")
+
+
+def start(key=''):
+    """
+    Start training models with Ultralytics HUB (DEPRECATED).
+
+    Args:
+        key (str, optional): A string containing either the API key and model ID combination (apikey_modelid),
+                               or the full model URL (https://hub.ultralytics.com/models/apikey_modelid).
+    """
+    api_key, model_id = key.split('_')
+    LOGGER.warning(f"""
+WARNING ⚠️ ultralytics.start() is deprecated after 8.0.60. Updated usage to train Ultralytics HUB models is:
+
+from ultralytics import YOLO, hub
+
+hub.login('{api_key}')
+model = YOLO('https://hub.ultralytics.com/models/{model_id}')
+model.train()""")
+
+
+def reset_model(model_id=''):
+    """Reset a trained model to an untrained state."""
+    r = requests.post('https://api.ultralytics.com/model-reset', json={'apiKey': Auth().api_key, 'modelId': model_id})
+    if r.status_code == 200:
+        LOGGER.info(f'{PREFIX}Model reset successfully')
+        return
+    LOGGER.warning(f'{PREFIX}Model reset failure {r.status_code} {r.reason}')
+
+
+def export_fmts_hub():
+    """Returns a list of HUB-supported export formats."""
+    from ultralytics.yolo.engine.exporter import export_formats
+    return list(export_formats()['Argument'][1:]) + ['ultralytics_tflite', 'ultralytics_coreml']
+
+
+def export_model(model_id='', format='torchscript'):
+    """Export a model to all formats."""
+    assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
+    r = requests.post(f'https://api.ultralytics.com/v1/models/{model_id}/export',
+                      json={'format': format},
+                      headers={'x-api-key': Auth().api_key})
+    assert r.status_code == 200, f'{PREFIX}{format} export failure {r.status_code} {r.reason}'
+    LOGGER.info(f'{PREFIX}{format} export started ✅')
+
+
+def get_export(model_id='', format='torchscript'):
+    """Get an exported model dictionary with download URL."""
+    assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
+    r = requests.post('https://api.ultralytics.com/get-export',
+                      json={
+                          'apiKey': Auth().api_key,
+                          'modelId': model_id,
+                          'format': format})
+    assert r.status_code == 200, f'{PREFIX}{format} get_export failure {r.status_code} {r.reason}'
+    return r.json()
+
+
+def check_dataset(path='', task='detect'):
+    """
+    Function for error-checking HUB dataset Zip file before upload. It checks a dataset for errors before it is
+    uploaded to the HUB. Usage examples are given below.
+
+    Args:
+        path (str, optional): Path to data.zip (with data.yaml inside data.zip). Defaults to ''.
+        task (str, optional): Dataset task. Options are 'detect', 'segment', 'pose', 'classify'. Defaults to 'detect'.
+
+    Example:
+        ```python
+        from ultralytics.hub import check_dataset
+
+        check_dataset('path/to/coco8.zip', task='detect')  # detect dataset
+        check_dataset('path/to/coco8-seg.zip', task='segment')  # segment dataset
+        check_dataset('path/to/coco8-pose.zip', task='pose')  # pose dataset
+        ```
+    """
+    HUBDatasetStats(path=path, task=task).get_json()
+    LOGGER.info('Checks completed correctly ✅. Upload this dataset to https://hub.ultralytics.com/datasets/.')
+
+
+if __name__ == '__main__':
+    start()

ultralytics/hub/auth.py

@@ -0,0 +1,139 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import requests
+
+from ultralytics.hub.utils import HUB_API_ROOT, PREFIX, request_with_credentials
+from ultralytics.yolo.utils import LOGGER, SETTINGS, emojis, is_colab, set_settings
+
+API_KEY_URL = 'https://hub.ultralytics.com/settings?tab=api+keys'
+
+
+class Auth:
+    id_token = api_key = model_key = False
+
+    def __init__(self, api_key='', verbose=False):
+        """
+        Initialize the Auth class with an optional API key.
+
+        Args:
+            api_key (str, optional): May be an API key or a combination API key and model ID, i.e. key_id
+        """
+        # Split the input API key in case it contains a combined key_model and keep only the API key part
+        api_key = api_key.split('_')[0]
+
+        # Set API key attribute as value passed or SETTINGS API key if none passed
+        self.api_key = api_key or SETTINGS.get('api_key', '')
+
+        # If an API key is provided
+        if self.api_key:
+            # If the provided API key matches the API key in the SETTINGS
+            if self.api_key == SETTINGS.get('api_key'):
+                # Log that the user is already logged in
+                if verbose:
+                    LOGGER.info(f'{PREFIX}Authenticated ✅')
+                return
+            else:
+                # Attempt to authenticate with the provided API key
+                success = self.authenticate()
+        # If the API key is not provided and the environment is a Google Colab notebook
+        elif is_colab():
+            # Attempt to authenticate using browser cookies
+            success = self.auth_with_cookies()
+        else:
+            # Request an API key
+            success = self.request_api_key()
+
+        # Update SETTINGS with the new API key after successful authentication
+        if success:
+            set_settings({'api_key': self.api_key})
+            # Log that the new login was successful
+            if verbose:
+                LOGGER.info(f'{PREFIX}New authentication successful ✅')
+        elif verbose:
+            LOGGER.info(f'{PREFIX}Retrieve API key from {API_KEY_URL}')
+
+    def request_api_key(self, max_attempts=3):
+        """
+        Prompt the user to input their API key. Returns the model ID.
+        """
+        import getpass
+        for attempts in range(max_attempts):
+            LOGGER.info(f'{PREFIX}Login. Attempt {attempts + 1} of {max_attempts}')
+            input_key = getpass.getpass(f'Enter API key from {API_KEY_URL} ')
+            self.api_key = input_key.split('_')[0]  # remove model id if present
+            if self.authenticate():
+                return True
+        raise ConnectionError(emojis(f'{PREFIX}Failed to authenticate ❌'))
+
+    def authenticate(self) -> bool:
+        """
+        Attempt to authenticate with the server using either id_token or API key.
+
+        Returns:
+            bool: True if authentication is successful, False otherwise.
+        """
+        try:
+            header = self.get_auth_header()
+            if header:
+                r = requests.post(f'{HUB_API_ROOT}/v1/auth', headers=header)
+                if not r.json().get('success', False):
+                    raise ConnectionError('Unable to authenticate.')
+                return True
+            raise ConnectionError('User has not authenticated locally.')
+        except ConnectionError:
+            self.id_token = self.api_key = False  # reset invalid
+            LOGGER.warning(f'{PREFIX}Invalid API key ⚠️')
+            return False
+
+    def auth_with_cookies(self) -> bool:
+        """
+        Attempt to fetch authentication via cookies and set id_token.
+        User must be logged in to HUB and running in a supported browser.
+
+        Returns:
+            bool: True if authentication is successful, False otherwise.
+        """
+        if not is_colab():
+            return False  # Currently only works with Colab
+        try:
+            authn = request_with_credentials(f'{HUB_API_ROOT}/v1/auth/auto')
+            if authn.get('success', False):
+                self.id_token = authn.get('data', {}).get('idToken', None)
+                self.authenticate()
+                return True
+            raise ConnectionError('Unable to fetch browser authentication details.')
+        except ConnectionError:
+            self.id_token = False  # reset invalid
+            return False
+
+    def get_auth_header(self):
+        """
+        Get the authentication header for making API requests.
+
+        Returns:
+            (dict): The authentication header if id_token or API key is set, None otherwise.
+        """
+        if self.id_token:
+            return {'authorization': f'Bearer {self.id_token}'}
+        elif self.api_key:
+            return {'x-api-key': self.api_key}
+        else:
+            return None
+
+    def get_state(self) -> bool:
+        """
+        Get the authentication state.
+
+        Returns:
+            bool: True if either id_token or API key is set, False otherwise.
+        """
+        return self.id_token or self.api_key
+
+    def set_api_key(self, key: str):
+        """
+        Set the API key for authentication.
+
+        Args:
+            key (str): The API key string.
+        """
+        self.api_key = key

ultralytics/hub/session.py

@@ -0,0 +1,189 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+import signal
+import sys
+from pathlib import Path
+from time import sleep
+
+import requests
+
+from ultralytics.hub.utils import HUB_API_ROOT, PREFIX, smart_request
+from ultralytics.yolo.utils import LOGGER, __version__, checks, emojis, is_colab, threaded
+from ultralytics.yolo.utils.errors import HUBModelError
+
+AGENT_NAME = f'python-{__version__}-colab' if is_colab() else f'python-{__version__}-local'
+
+
+class HUBTrainingSession:
+    """
+    HUB training session for Ultralytics HUB YOLO models. Handles model initialization, heartbeats, and checkpointing.
+
+    Args:
+        url (str): Model identifier used to initialize the HUB training session.
+
+    Attributes:
+        agent_id (str): Identifier for the instance communicating with the server.
+        model_id (str): Identifier for the YOLOv5 model being trained.
+        model_url (str): URL for the model in Ultralytics HUB.
+        api_url (str): API URL for the model in Ultralytics HUB.
+        auth_header (Dict): Authentication header for the Ultralytics HUB API requests.
+        rate_limits (Dict): Rate limits for different API calls (in seconds).
+        timers (Dict): Timers for rate limiting.
+        metrics_queue (Dict): Queue for the model's metrics.
+        model (Dict): Model data fetched from Ultralytics HUB.
+        alive (bool): Indicates if the heartbeat loop is active.
+    """
+
+    def __init__(self, url):
+        """
+        Initialize the HUBTrainingSession with the provided model identifier.
+
+        Args:
+            url (str): Model identifier used to initialize the HUB training session.
+                         It can be a URL string or a model key with specific format.
+
+        Raises:
+            ValueError: If the provided model identifier is invalid.
+            ConnectionError: If connecting with global API key is not supported.
+        """
+
+        from ultralytics.hub.auth import Auth
+
+        # Parse input
+        if url.startswith('https://hub.ultralytics.com/models/'):
+            url = url.split('https://hub.ultralytics.com/models/')[-1]
+        if [len(x) for x in url.split('_')] == [42, 20]:
+            key, model_id = url.split('_')
+        elif len(url) == 20:
+            key, model_id = '', url
+        else:
+            raise HUBModelError(f"model='{url}' not found. Check format is correct, i.e. "
+                                f"model='https://hub.ultralytics.com/models/MODEL_ID' and try again.")
+
+        # Authorize
+        auth = Auth(key)
+        self.agent_id = None  # identifies which instance is communicating with server
+        self.model_id = model_id
+        self.model_url = f'https://hub.ultralytics.com/models/{model_id}'
+        self.api_url = f'{HUB_API_ROOT}/v1/models/{model_id}'
+        self.auth_header = auth.get_auth_header()
+        self.rate_limits = {'metrics': 3.0, 'ckpt': 900.0, 'heartbeat': 300.0}  # rate limits (seconds)
+        self.timers = {}  # rate limit timers (seconds)
+        self.metrics_queue = {}  # metrics queue
+        self.model = self._get_model()
+        self.alive = True
+        self._start_heartbeat()  # start heartbeats
+        self._register_signal_handlers()
+        LOGGER.info(f'{PREFIX}View model at {self.model_url} 🚀')
+
+    def _register_signal_handlers(self):
+        """Register signal handlers for SIGTERM and SIGINT signals to gracefully handle termination."""
+        signal.signal(signal.SIGTERM, self._handle_signal)
+        signal.signal(signal.SIGINT, self._handle_signal)
+
+    def _handle_signal(self, signum, frame):
+        """
+        Handle kill signals and prevent heartbeats from being sent on Colab after termination.
+        This method does not use frame, it is included as it is passed by signal.
+        """
+        if self.alive is True:
+            LOGGER.info(f'{PREFIX}Kill signal received! ❌')
+            self._stop_heartbeat()
+            sys.exit(signum)
+
+    def _stop_heartbeat(self):
+        """Terminate the heartbeat loop."""
+        self.alive = False
+
+    def upload_metrics(self):
+        """Upload model metrics to Ultralytics HUB."""
+        payload = {'metrics': self.metrics_queue.copy(), 'type': 'metrics'}
+        smart_request('post', self.api_url, json=payload, headers=self.auth_header, code=2)
+
+    def _get_model(self):
+        """Fetch and return model data from Ultralytics HUB."""
+        api_url = f'{HUB_API_ROOT}/v1/models/{self.model_id}'
+
+        try:
+            response = smart_request('get', api_url, headers=self.auth_header, thread=False, code=0)
+            data = response.json().get('data', None)
+
+            if data.get('status', None) == 'trained':
+                raise ValueError(emojis(f'Model is already trained and uploaded to {self.model_url} 🚀'))
+
+            if not data.get('data', None):
+                raise ValueError('Dataset may still be processing. Please wait a minute and try again.')  # RF fix
+            self.model_id = data['id']
+
+            if data['status'] == 'new':  # new model to start training
+                self.train_args = {
+                    # TODO: deprecate 'batch_size' key for 'batch' in 3Q23
+                    'batch': data['batch' if ('batch' in data) else 'batch_size'],
+                    'epochs': data['epochs'],
+                    'imgsz': data['imgsz'],
+                    'patience': data['patience'],
+                    'device': data['device'],
+                    'cache': data['cache'],
+                    'data': data['data']}
+                self.model_file = data.get('cfg') or data.get('weights')  # cfg for pretrained=False
+                self.model_file = checks.check_yolov5u_filename(self.model_file, verbose=False)  # YOLOv5->YOLOv5u
+            elif data['status'] == 'training':  # existing model to resume training
+                self.train_args = {'data': data['data'], 'resume': True}
+                self.model_file = data['resume']
+
+            return data
+        except requests.exceptions.ConnectionError as e:
+            raise ConnectionRefusedError('ERROR: The HUB server is not online. Please try again later.') from e
+        except Exception:
+            raise
+
+    def upload_model(self, epoch, weights, is_best=False, map=0.0, final=False):
+        """
+        Upload a model checkpoint to Ultralytics HUB.
+
+        Args:
+            epoch (int): The current training epoch.
+            weights (str): Path to the model weights file.
+            is_best (bool): Indicates if the current model is the best one so far.
+            map (float): Mean average precision of the model.
+            final (bool): Indicates if the model is the final model after training.
+        """
+        if Path(weights).is_file():
+            with open(weights, 'rb') as f:
+                file = f.read()
+        else:
+            LOGGER.warning(f'{PREFIX}WARNING ⚠️ Model upload issue. Missing model {weights}.')
+            file = None
+        url = f'{self.api_url}/upload'
+        # url = 'http://httpbin.org/post'  # for debug
+        data = {'epoch': epoch}
+        if final:
+            data.update({'type': 'final', 'map': map})
+            smart_request('post',
+                          url,
+                          data=data,
+                          files={'best.pt': file},
+                          headers=self.auth_header,
+                          retry=10,
+                          timeout=3600,
+                          thread=False,
+                          progress=True,
+                          code=4)
+        else:
+            data.update({'type': 'epoch', 'isBest': bool(is_best)})
+            smart_request('post', url, data=data, files={'last.pt': file}, headers=self.auth_header, code=3)
+
+    @threaded
+    def _start_heartbeat(self):
+        """Begin a threaded heartbeat loop to report the agent's status to Ultralytics HUB."""
+        while self.alive:
+            r = smart_request('post',
+                              f'{HUB_API_ROOT}/v1/agent/heartbeat/models/{self.model_id}',
+                              json={
+                                  'agent': AGENT_NAME,
+                                  'agentId': self.agent_id},
+                              headers=self.auth_header,
+                              retry=0,
+                              code=5,
+                              thread=False)  # already in a thread
+            self.agent_id = r.json().get('data', {}).get('agentId', None)
+            sleep(self.rate_limits['heartbeat'])

ultralytics/hub/utils.py

@@ -0,0 +1,217 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import os
+import platform
+import random
+import sys
+import threading
+import time
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+from ultralytics.yolo.utils import (ENVIRONMENT, LOGGER, ONLINE, RANK, SETTINGS, TESTS_RUNNING, TQDM_BAR_FORMAT,
+                                    TryExcept, __version__, colorstr, get_git_origin_url, is_colab, is_git_dir,
+                                    is_pip_package)
+
+PREFIX = colorstr('Ultralytics HUB: ')
+HELP_MSG = 'If this issue persists please visit https://github.com/ultralytics/hub/issues for assistance.'
+HUB_API_ROOT = os.environ.get('ULTRALYTICS_HUB_API', 'https://api.ultralytics.com')
+
+
+def request_with_credentials(url: str) -> any:
+    """
+    Make an AJAX request with cookies attached in a Google Colab environment.
+
+    Args:
+        url (str): The URL to make the request to.
+
+    Returns:
+        (any): The response data from the AJAX request.
+
+    Raises:
+        OSError: If the function is not run in a Google Colab environment.
+    """
+    if not is_colab():
+        raise OSError('request_with_credentials() must run in a Colab environment')
+    from google.colab import output  # noqa
+    from IPython import display  # noqa
+    display.display(
+        display.Javascript("""
+            window._hub_tmp = new Promise((resolve, reject) => {
+                const timeout = setTimeout(() => reject("Failed authenticating existing browser session"), 5000)
+                fetch("%s", {
+                    method: 'POST',
+                    credentials: 'include'
+                })
+                    .then((response) => resolve(response.json()))
+                    .then((json) => {
+                    clearTimeout(timeout);
+                    }).catch((err) => {
+                    clearTimeout(timeout);
+                    reject(err);
+                });
+            });
+            """ % url))
+    return output.eval_js('_hub_tmp')
+
+
+def requests_with_progress(method, url, **kwargs):
+    """
+    Make an HTTP request using the specified method and URL, with an optional progress bar.
+
+    Args:
+        method (str): The HTTP method to use (e.g. 'GET', 'POST').
+        url (str): The URL to send the request to.
+        **kwargs (dict): Additional keyword arguments to pass to the underlying `requests.request` function.
+
+    Returns:
+        (requests.Response): The response object from the HTTP request.
+
+    Note:
+        If 'progress' is set to True, the progress bar will display the download progress
+        for responses with a known content length.
+    """
+    progress = kwargs.pop('progress', False)
+    if not progress:
+        return requests.request(method, url, **kwargs)
+    response = requests.request(method, url, stream=True, **kwargs)
+    total = int(response.headers.get('content-length', 0))  # total size
+    pbar = tqdm(total=total, unit='B', unit_scale=True, unit_divisor=1024, bar_format=TQDM_BAR_FORMAT)
+    for data in response.iter_content(chunk_size=1024):
+        pbar.update(len(data))
+    pbar.close()
+    return response
+
+
+def smart_request(method, url, retry=3, timeout=30, thread=True, code=-1, verbose=True, progress=False, **kwargs):
+    """
+    Makes an HTTP request using the 'requests' library, with exponential backoff retries up to a specified timeout.
+
+    Args:
+        method (str): The HTTP method to use for the request. Choices are 'post' and 'get'.
+        url (str): The URL to make the request to.
+        retry (int, optional): Number of retries to attempt before giving up. Default is 3.
+        timeout (int, optional): Timeout in seconds after which the function will give up retrying. Default is 30.
+        thread (bool, optional): Whether to execute the request in a separate daemon thread. Default is True.
+        code (int, optional): An identifier for the request, used for logging purposes. Default is -1.
+        verbose (bool, optional): A flag to determine whether to print out to console or not. Default is True.
+        progress (bool, optional): Whether to show a progress bar during the request. Default is False.
+        **kwargs (dict): Keyword arguments to be passed to the requests function specified in method.
+
+    Returns:
+        (requests.Response): The HTTP response object. If the request is executed in a separate thread, returns None.
+    """
+    retry_codes = (408, 500)  # retry only these codes
+
+    @TryExcept(verbose=verbose)
+    def func(func_method, func_url, **func_kwargs):
+        """Make HTTP requests with retries and timeouts, with optional progress tracking."""
+        r = None  # response
+        t0 = time.time()  # initial time for timer
+        for i in range(retry + 1):
+            if (time.time() - t0) > timeout:
+                break
+            r = requests_with_progress(func_method, func_url, **func_kwargs)  # i.e. get(url, data, json, files)
+            if r.status_code < 300:  # return codes in the 2xx range are generally considered "good" or "successful"
+                break
+            try:
+                m = r.json().get('message', 'No JSON message.')
+            except AttributeError:
+                m = 'Unable to read JSON.'
+            if i == 0:
+                if r.status_code in retry_codes:
+                    m += f' Retrying {retry}x for {timeout}s.' if retry else ''
+                elif r.status_code == 429:  # rate limit
+                    h = r.headers  # response headers
+                    m = f"Rate limit reached ({h['X-RateLimit-Remaining']}/{h['X-RateLimit-Limit']}). " \
+                        f"Please retry after {h['Retry-After']}s."
+                if verbose:
+                    LOGGER.warning(f'{PREFIX}{m} {HELP_MSG} ({r.status_code} #{code})')
+                if r.status_code not in retry_codes:
+                    return r
+            time.sleep(2 ** i)  # exponential standoff
+        return r
+
+    args = method, url
+    kwargs['progress'] = progress
+    if thread:
+        threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True).start()
+    else:
+        return func(*args, **kwargs)
+
+
+class Events:
+    """
+    A class for collecting anonymous event analytics. Event analytics are enabled when sync=True in settings and
+    disabled when sync=False. Run 'yolo settings' to see and update settings YAML file.
+
+    Attributes:
+        url (str): The URL to send anonymous events.
+        rate_limit (float): The rate limit in seconds for sending events.
+        metadata (dict): A dictionary containing metadata about the environment.
+        enabled (bool): A flag to enable or disable Events based on certain conditions.
+    """
+
+    url = 'https://www.google-analytics.com/mp/collect?measurement_id=G-X8NCJYTQXM&api_secret=QLQrATrNSwGRFRLE-cbHJw'
+
+    def __init__(self):
+        """
+        Initializes the Events object with default values for events, rate_limit, and metadata.
+        """
+        self.events = []  # events list
+        self.rate_limit = 60.0  # rate limit (seconds)
+        self.t = 0.0  # rate limit timer (seconds)
+        self.metadata = {
+            'cli': Path(sys.argv[0]).name == 'yolo',
+            'install': 'git' if is_git_dir() else 'pip' if is_pip_package() else 'other',
+            'python': '.'.join(platform.python_version_tuple()[:2]),  # i.e. 3.10
+            'version': __version__,
+            'env': ENVIRONMENT,
+            'session_id': round(random.random() * 1E15),
+            'engagement_time_msec': 1000}
+        self.enabled = \
+            SETTINGS['sync'] and \
+            RANK in (-1, 0) and \
+            not TESTS_RUNNING and \
+            ONLINE and \
+            (is_pip_package() or get_git_origin_url() == 'https://github.com/ultralytics/ultralytics.git')
+
+    def __call__(self, cfg):
+        """
+        Attempts to add a new event to the events list and send events if the rate limit is reached.
+
+        Args:
+            cfg (IterableSimpleNamespace): The configuration object containing mode and task information.
+        """
+        if not self.enabled:
+            # Events disabled, do nothing
+            return
+
+        # Attempt to add to events
+        if len(self.events) < 25:  # Events list limited to 25 events (drop any events past this)
+            params = {**self.metadata, **{'task': cfg.task}}
+            if cfg.mode == 'export':
+                params['format'] = cfg.format
+            self.events.append({'name': cfg.mode, 'params': params})
+
+        # Check rate limit
+        t = time.time()
+        if (t - self.t) < self.rate_limit:
+            # Time is under rate limiter, wait to send
+            return
+
+        # Time is over rate limiter, send now
+        data = {'client_id': SETTINGS['uuid'], 'events': self.events}  # SHA-256 anonymized UUID hash and events list
+
+        # POST equivalent to requests.post(self.url, json=data)
+        smart_request('post', self.url, json=data, retry=0, verbose=False)
+
+        # Reset events and rate limit timer
+        self.events = []
+        self.t = t
+
+
+# Run below code on hub/utils init -------------------------------------------------------------------------------------
+events = Events()

ultralytics/models/README.md

@@ -0,0 +1,45 @@
+## Models
+
+Welcome to the Ultralytics Models directory! Here you will find a wide variety of pre-configured model configuration
+files (`*.yaml`s) that can be used to create custom YOLO models. The models in this directory have been expertly crafted
+and fine-tuned by the Ultralytics team to provide the best performance for a wide range of object detection and image
+segmentation tasks.
+
+These model configurations cover a wide range of scenarios, from simple object detection to more complex tasks like
+instance segmentation and object tracking. They are also designed to run efficiently on a variety of hardware platforms,
+from CPUs to GPUs. Whether you are a seasoned machine learning practitioner or just getting started with YOLO, this
+directory provides a great starting point for your custom model development needs.
+
+To get started, simply browse through the models in this directory and find one that best suits your needs. Once you've
+selected a model, you can use the provided `*.yaml` file to train and deploy your custom YOLO model with ease. See full
+details at the Ultralytics [Docs](https://docs.ultralytics.com/models), and if you need help or have any questions, feel free
+to reach out to the Ultralytics team for support. So, don't wait, start creating your custom YOLO model now!
+
+### Usage
+
+Model `*.yaml` files may be used directly in the Command Line Interface (CLI) with a `yolo` command:
+
+```bash
+yolo task=detect mode=train model=yolov8n.yaml data=coco128.yaml epochs=100
+```
+
+They may also be used directly in a Python environment, and accepts the same
+[arguments](https://docs.ultralytics.com/usage/cfg/) as in the CLI example above:
+
+```python
+from ultralytics import YOLO
+
+model = YOLO("model.yaml")  # build a YOLOv8n model from scratch
+# YOLO("model.pt")  use pre-trained model if available
+model.info()  # display model information
+model.train(data="coco128.yaml", epochs=100)  # train the model
+```
+
+## Pre-trained Model Architectures
+
+Ultralytics supports many model architectures. Visit https://docs.ultralytics.com/models to view detailed information
+and usage. Any of these models can be used by loading their configs or pretrained checkpoints if available.
+
+## Contributing New Models
+
+If you've developed a new model architecture or have improvements for existing models that you'd like to contribute to the Ultralytics community, please submit your contribution in a new Pull Request. For more details, visit our [Contributing Guide](https://docs.ultralytics.com/help/contributing).

ultralytics/models/rt-detr/rtdetr-l.yaml

@@ -0,0 +1,50 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# RT-DETR-l object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/rtdetr
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n'
+  # [depth, width, max_channels]
+  l: [1.00, 1.00, 1024]
+
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, HGStem, [32, 48]]  # 0-P2/4
+  - [-1, 6, HGBlock, [48, 128, 3]]  # stage 1
+
+  - [-1, 1, DWConv, [128, 3, 2, 1, False]]  # 2-P3/8
+  - [-1, 6, HGBlock, [96, 512, 3]]   # stage 2
+
+  - [-1, 1, DWConv, [512, 3, 2, 1, False]]  # 4-P3/16
+  - [-1, 6, HGBlock, [192, 1024, 5, True, False]]  # cm, c2, k, light, shortcut
+  - [-1, 6, HGBlock, [192, 1024, 5, True, True]]
+  - [-1, 6, HGBlock, [192, 1024, 5, True, True]]  # stage 3
+
+  - [-1, 1, DWConv, [1024, 3, 2, 1, False]]  # 8-P4/32
+  - [-1, 6, HGBlock, [384, 2048, 5, True, False]]  # stage 4
+
+head:
+  - [-1, 1, Conv, [256, 1, 1, None, 1, 1, False]]  # 10 input_proj.2
+  - [-1, 1, AIFI, [1024, 8]]
+  - [-1, 1, Conv, [256, 1, 1]]   # 12, Y5, lateral_convs.0
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [7, 1, Conv, [256, 1, 1, None, 1, 1, False]]  # 14 input_proj.1
+  - [[-2, -1], 1, Concat, [1]]
+  - [-1, 3, RepC3, [256]]  # 16, fpn_blocks.0
+  - [-1, 1, Conv, [256, 1, 1]]   # 17, Y4, lateral_convs.1
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [3, 1, Conv, [256, 1, 1, None, 1, 1, False]]  # 19 input_proj.0
+  - [[-2, -1], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, RepC3, [256]]    # X3 (21), fpn_blocks.1
+
+  - [-1, 1, Conv, [256, 3, 2]]   # 22, downsample_convs.0
+  - [[-1, 17], 1, Concat, [1]]  # cat Y4
+  - [-1, 3, RepC3, [256]]    # F4 (24), pan_blocks.0
+
+  - [-1, 1, Conv, [256, 3, 2]]   # 25, downsample_convs.1
+  - [[-1, 12], 1, Concat, [1]]  # cat Y5
+  - [-1, 3, RepC3, [256]]    # F5 (27), pan_blocks.1
+
+  - [[21, 24, 27], 1, RTDETRDecoder, [nc]]  # Detect(P3, P4, P5)

ultralytics/models/rt-detr/rtdetr-x.yaml

@@ -0,0 +1,54 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# RT-DETR-x object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/rtdetr
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n'
+  # [depth, width, max_channels]
+  x: [1.00, 1.00, 2048]
+
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, HGStem, [32, 64]]  # 0-P2/4
+  - [-1, 6, HGBlock, [64, 128, 3]]  # stage 1
+
+  - [-1, 1, DWConv, [128, 3, 2, 1, False]]  # 2-P3/8
+  - [-1, 6, HGBlock, [128, 512, 3]]
+  - [-1, 6, HGBlock, [128, 512, 3, False, True]]   # 4-stage 2
+
+  - [-1, 1, DWConv, [512, 3, 2, 1, False]]  # 5-P3/16
+  - [-1, 6, HGBlock, [256, 1024, 5, True, False]]  # cm, c2, k, light, shortcut
+  - [-1, 6, HGBlock, [256, 1024, 5, True, True]]
+  - [-1, 6, HGBlock, [256, 1024, 5, True, True]]
+  - [-1, 6, HGBlock, [256, 1024, 5, True, True]]
+  - [-1, 6, HGBlock, [256, 1024, 5, True, True]]  # 10-stage 3
+
+  - [-1, 1, DWConv, [1024, 3, 2, 1, False]]  # 11-P4/32
+  - [-1, 6, HGBlock, [512, 2048, 5, True, False]]
+  - [-1, 6, HGBlock, [512, 2048, 5, True, True]]  # 13-stage 4
+
+head:
+  - [-1, 1, Conv, [384, 1, 1, None, 1, 1, False]]  # 14 input_proj.2
+  - [-1, 1, AIFI, [2048, 8]]
+  - [-1, 1, Conv, [384, 1, 1]]   # 16, Y5, lateral_convs.0
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [10, 1, Conv, [384, 1, 1, None, 1, 1, False]]  # 18 input_proj.1
+  - [[-2, -1], 1, Concat, [1]]
+  - [-1, 3, RepC3, [384]]  # 20, fpn_blocks.0
+  - [-1, 1, Conv, [384, 1, 1]]   # 21, Y4, lateral_convs.1
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [4, 1, Conv, [384, 1, 1, None, 1, 1, False]]  # 23 input_proj.0
+  - [[-2, -1], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, RepC3, [384]]    # X3 (25), fpn_blocks.1
+
+  - [-1, 1, Conv, [384, 3, 2]]   # 26, downsample_convs.0
+  - [[-1, 21], 1, Concat, [1]]  # cat Y4
+  - [-1, 3, RepC3, [384]]    # F4 (28), pan_blocks.0
+
+  - [-1, 1, Conv, [384, 3, 2]]   # 29, downsample_convs.1
+  - [[-1, 16], 1, Concat, [1]]  # cat Y5
+  - [-1, 3, RepC3, [384]]    # F5 (31), pan_blocks.1
+
+  - [[25, 28, 31], 1, RTDETRDecoder, [nc]]  # Detect(P3, P4, P5)

ultralytics/models/v3/yolov3-spp.yaml

@@ -0,0 +1,48 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv3-SPP object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/yolov3
+
+# Parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# darknet53 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [32, 3, 1]],  # 0
+   [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
+   [-1, 1, Bottleneck, [64]],
+   [-1, 1, Conv, [128, 3, 2]],  # 3-P2/4
+   [-1, 2, Bottleneck, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 5-P3/8
+   [-1, 8, Bottleneck, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 7-P4/16
+   [-1, 8, Bottleneck, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 9-P5/32
+   [-1, 4, Bottleneck, [1024]],  # 10
+  ]
+
+# YOLOv3-SPP head
+head:
+  [[-1, 1, Bottleneck, [1024, False]],
+   [-1, 1, SPP, [512, [5, 9, 13]]],
+   [-1, 1, Conv, [1024, 3, 1]],
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, Conv, [1024, 3, 1]],  # 15 (P5/32-large)
+
+   [-2, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, Conv, [512, 3, 1]],  # 22 (P4/16-medium)
+
+   [-2, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P3
+   [-1, 1, Bottleneck, [256, False]],
+   [-1, 2, Bottleneck, [256, False]],  # 27 (P3/8-small)
+
+   [[27, 22, 15], 1, Detect, [nc]],   # Detect(P3, P4, P5)
+  ]

ultralytics/models/v3/yolov3-tiny.yaml

@@ -0,0 +1,39 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv3-tiny object detection model with P4-P5 outputs. For details see https://docs.ultralytics.com/models/yolov3
+
+# Parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# YOLOv3-tiny backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [16, 3, 1]],  # 0
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 1-P1/2
+   [-1, 1, Conv, [32, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 3-P2/4
+   [-1, 1, Conv, [64, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 5-P3/8
+   [-1, 1, Conv, [128, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 7-P4/16
+   [-1, 1, Conv, [256, 3, 1]],
+   [-1, 1, nn.MaxPool2d, [2, 2, 0]],  # 9-P5/32
+   [-1, 1, Conv, [512, 3, 1]],
+   [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]],  # 11
+   [-1, 1, nn.MaxPool2d, [2, 1, 0]],  # 12
+  ]
+
+# YOLOv3-tiny head
+head:
+  [[-1, 1, Conv, [1024, 3, 1]],
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, Conv, [512, 3, 1]],  # 15 (P5/32-large)
+
+   [-2, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Conv, [256, 3, 1]],  # 19 (P4/16-medium)
+
+   [[19, 15], 1, Detect, [nc]],  # Detect(P4, P5)
+  ]

ultralytics/models/v3/yolov3.yaml

@@ -0,0 +1,48 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv3 object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/yolov3
+
+# Parameters
+nc: 80  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# darknet53 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [32, 3, 1]],  # 0
+   [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
+   [-1, 1, Bottleneck, [64]],
+   [-1, 1, Conv, [128, 3, 2]],  # 3-P2/4
+   [-1, 2, Bottleneck, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 5-P3/8
+   [-1, 8, Bottleneck, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 7-P4/16
+   [-1, 8, Bottleneck, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 9-P5/32
+   [-1, 4, Bottleneck, [1024]],  # 10
+  ]
+
+# YOLOv3 head
+head:
+  [[-1, 1, Bottleneck, [1024, False]],
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, Conv, [1024, 3, 1]],
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, Conv, [1024, 3, 1]],  # 15 (P5/32-large)
+
+   [-2, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Bottleneck, [512, False]],
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, Conv, [512, 3, 1]],  # 22 (P4/16-medium)
+
+   [-2, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P3
+   [-1, 1, Bottleneck, [256, False]],
+   [-1, 2, Bottleneck, [256, False]],  # 27 (P3/8-small)
+
+   [[27, 22, 15], 1, Detect, [nc]],   # Detect(P3, P4, P5)
+  ]

ultralytics/models/v5/yolov5-p6.yaml

@@ -0,0 +1,61 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv5 object detection model with P3-P6 outputs. For details see https://docs.ultralytics.com/models/yolov5
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov5n-p6.yaml' will call yolov5-p6.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 1024]
+  l: [1.00, 1.00, 1024]
+  x: [1.33, 1.25, 1024]
+
+# YOLOv5 v6.0 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 6, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [768, 3, 2]],  # 7-P5/32
+   [-1, 3, C3, [768]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 9-P6/64
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 11
+  ]
+
+# YOLOv5 v6.0 head
+head:
+  [[-1, 1, Conv, [768, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 8], 1, Concat, [1]],  # cat backbone P5
+   [-1, 3, C3, [768, False]],  # 15
+
+   [-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 19
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 23 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 20], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 26 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 16], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [768, False]],  # 29 (P5/32-large)
+
+   [-1, 1, Conv, [768, 3, 2]],
+   [[-1, 12], 1, Concat, [1]],  # cat head P6
+   [-1, 3, C3, [1024, False]],  # 32 (P6/64-xlarge)
+
+   [[23, 26, 29, 32], 1, Detect, [nc]],  # Detect(P3, P4, P5, P6)
+  ]

ultralytics/models/v5/yolov5.yaml

@@ -0,0 +1,50 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv5 object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/yolov5
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov5n.yaml' will call yolov5.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 1024]
+  l: [1.00, 1.00, 1024]
+  x: [1.33, 1.25, 1024]
+
+# YOLOv5 v6.0 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 6, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
+  ]
+
+# YOLOv5 v6.0 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
+
+   [[17, 20, 23], 1, Detect, [nc]],  # Detect(P3, P4, P5)
+  ]

ultralytics/models/v6/yolov6.yaml

@@ -0,0 +1,53 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv6 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/models/yolov6
+
+# Parameters
+nc: 80  # number of classes
+activation: nn.ReLU()  # (optional) model default activation function
+scales: # model compound scaling constants, i.e. 'model=yolov6n.yaml' will call yolov8.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 768]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.25, 512]
+
+# YOLOv6-3.0s backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 6, Conv, [128, 3, 1]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 12, Conv, [256, 3, 1]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 18, Conv, [512, 3, 1]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 6, Conv, [1024, 3, 1]]
+  - [-1, 1, SPPF, [1024, 5]]  # 9
+
+# YOLOv6-3.0s head
+head:
+  - [-1, 1, Conv, [256, 1, 1]]
+  - [-1, 1, nn.ConvTranspose2d, [256, 2, 2, 0]]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 1, Conv, [256, 3, 1]]
+  - [-1, 9, Conv, [256, 3, 1]]  # 14
+
+  - [-1, 1, Conv, [128, 1, 1]]
+  - [-1, 1, nn.ConvTranspose2d, [128, 2, 2, 0]]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 1, Conv, [128, 3, 1]]
+  - [-1, 9, Conv, [128, 3, 1]]  # 19
+
+  - [-1, 1, Conv, [128, 3, 2]]
+  - [[-1, 15], 1, Concat, [1]]  # cat head P4
+  - [-1, 1, Conv, [256, 3, 1]]
+  - [-1, 9, Conv, [256, 3, 1]]  # 23
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 10], 1, Concat, [1]]  # cat head P5
+  - [-1, 1, Conv, [512, 3, 1]]
+  - [-1, 9, Conv, [512, 3, 1]]  # 27
+
+  - [[19, 23, 27], 1, Detect, [nc]]  # Detect(P3, P4, P5)

ultralytics/models/v8/yolov8-cls.yaml

@@ -0,0 +1,29 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8-cls image classification model. For Usage examples see https://docs.ultralytics.com/tasks/classify
+
+# Parameters
+nc: 1000  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 1024]
+  l: [1.00, 1.00, 1024]
+  x: [1.00, 1.25, 1024]
+
+# YOLOv8.0n backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+
+# YOLOv8.0n head
+head:
+  - [-1, 1, Classify, [nc]]  # Classify

ultralytics/models/v8/yolov8-p2.yaml

@@ -0,0 +1,54 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8 object detection model with P2-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 768]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.25, 512]
+
+# YOLOv8.0 backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 9
+
+# YOLOv8.0-p2 head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2f, [512]]  # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 2], 1, Concat, [1]]  # cat backbone P2
+  - [-1, 3, C2f, [128]]  # 18 (P2/4-xsmall)
+
+  - [-1, 1, Conv, [128, 3, 2]]
+  - [[-1, 15], 1, Concat, [1]]  # cat head P3
+  - [-1, 3, C2f, [256]]  # 21 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 12], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2f, [512]]  # 24 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 9], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2f, [1024]]  # 27 (P5/32-large)
+
+  - [[18, 21, 24, 27], 1, Detect, [nc]]  # Detect(P2, P3, P4, P5)

ultralytics/models/v8/yolov8-p6.yaml

@@ -0,0 +1,56 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8 object detection model with P3-P6 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n-p6.yaml' will call yolov8-p6.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 768]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.25, 512]
+
+# YOLOv8.0x6 backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [768, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [768, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 9-P6/64
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 11
+
+# YOLOv8.0x6 head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 8], 1, Concat, [1]]  # cat backbone P5
+  - [-1, 3, C2, [768, False]]  # 14
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2, [512, False]]  # 17
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2, [256, False]]  # 20 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 17], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2, [512, False]]  # 23 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 14], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2, [768, False]]  # 26 (P5/32-large)
+
+  - [-1, 1, Conv, [768, 3, 2]]
+  - [[-1, 11], 1, Concat, [1]]  # cat head P6
+  - [-1, 3, C2, [1024, False]]  # 29 (P6/64-xlarge)
+
+  - [[20, 23, 26, 29], 1, Detect, [nc]]  # Detect(P3, P4, P5, P6)

ultralytics/models/v8/yolov8-pose-p6.yaml

@@ -0,0 +1,57 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8-pose keypoints/pose estimation model. For Usage examples see https://docs.ultralytics.com/tasks/pose
+
+# Parameters
+nc: 1  # number of classes
+kpt_shape: [17, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
+scales: # model compound scaling constants, i.e. 'model=yolov8n-p6.yaml' will call yolov8-p6.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 768]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.25, 512]
+
+# YOLOv8.0x6 backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [768, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [768, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 9-P6/64
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 11
+
+# YOLOv8.0x6 head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 8], 1, Concat, [1]]  # cat backbone P5
+  - [-1, 3, C2, [768, False]]  # 14
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2, [512, False]]  # 17
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2, [256, False]]  # 20 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 17], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2, [512, False]]  # 23 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 14], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2, [768, False]]  # 26 (P5/32-large)
+
+  - [-1, 1, Conv, [768, 3, 2]]
+  - [[-1, 11], 1, Concat, [1]]  # cat head P6
+  - [-1, 3, C2, [1024, False]]  # 29 (P6/64-xlarge)
+
+  - [[20, 23, 26, 29], 1, Pose, [nc, kpt_shape]]  # Pose(P3, P4, P5, P6)

ultralytics/models/v8/yolov8-pose.yaml

@@ -0,0 +1,47 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8-pose keypoints/pose estimation model. For Usage examples see https://docs.ultralytics.com/tasks/pose
+
+# Parameters
+nc: 1  # number of classes
+kpt_shape: [17, 3]  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
+scales: # model compound scaling constants, i.e. 'model=yolov8n-pose.yaml' will call yolov8-pose.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 768]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.25, 512]
+
+# YOLOv8.0n backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 9
+
+# YOLOv8.0n head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2f, [512]]  # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 12], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 9], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, Pose, [nc, kpt_shape]]  # Pose(P3, P4, P5)

ultralytics/models/v8/yolov8-rtdetr.yaml

@@ -0,0 +1,46 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]  # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
+  s: [0.33, 0.50, 1024]  # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
+  m: [0.67, 0.75, 768]   # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
+  l: [1.00, 1.00, 512]   # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
+  x: [1.00, 1.25, 512]   # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
+
+# YOLOv8.0n backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 9
+
+# YOLOv8.0n head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2f, [512]]  # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 12], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 9], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, RTDETRDecoder, [nc]]  # Detect(P3, P4, P5)

ultralytics/models/v8/yolov8-seg.yaml

@@ -0,0 +1,46 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8-seg instance segmentation model. For Usage examples see https://docs.ultralytics.com/tasks/segment
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n-seg.yaml' will call yolov8-seg.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]
+  s: [0.33, 0.50, 1024]
+  m: [0.67, 0.75, 768]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.25, 512]
+
+# YOLOv8.0n backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 9
+
+# YOLOv8.0n head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2f, [512]]  # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 12], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 9], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, Segment, [nc, 32, 256]]  # Segment(P3, P4, P5)

ultralytics/models/v8/yolov8.yaml

@@ -0,0 +1,46 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
+
+# Parameters
+nc: 80  # number of classes
+scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
+  # [depth, width, max_channels]
+  n: [0.33, 0.25, 1024]  # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
+  s: [0.33, 0.50, 1024]  # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
+  m: [0.67, 0.75, 768]   # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
+  l: [1.00, 1.00, 512]   # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
+  x: [1.00, 1.25, 512]   # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
+
+# YOLOv8.0n backbone
+backbone:
+  # [from, repeats, module, args]
+  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
+  - [-1, 3, C2f, [128, True]]
+  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
+  - [-1, 6, C2f, [256, True]]
+  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
+  - [-1, 6, C2f, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
+  - [-1, 3, C2f, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5]]  # 9
+
+# YOLOv8.0n head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
+  - [-1, 3, C2f, [512]]  # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
+  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
+  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
+
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [[-1, 12], 1, Concat, [1]]  # cat head P4
+  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)
+
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [[-1, 9], 1, Concat, [1]]  # cat head P5
+  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, Detect, [nc]]  # Detect(P3, P4, P5)

ultralytics/nn/__init__.py

@@ -0,0 +1,9 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from .tasks import (BaseModel, ClassificationModel, DetectionModel, SegmentationModel, attempt_load_one_weight,
+                    attempt_load_weights, guess_model_scale, guess_model_task, parse_model, torch_safe_load,
+                    yaml_model_load)
+
+__all__ = ('attempt_load_one_weight', 'attempt_load_weights', 'parse_model', 'yaml_model_load', 'guess_model_task',
+           'guess_model_scale', 'torch_safe_load', 'DetectionModel', 'SegmentationModel', 'ClassificationModel',
+           'BaseModel')

ultralytics/nn/autobackend.py

@@ -0,0 +1,455 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import ast
+import contextlib
+import json
+import platform
+import zipfile
+from collections import OrderedDict, namedtuple
+from pathlib import Path
+from urllib.parse import urlparse
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+from ultralytics.yolo.utils import LINUX, LOGGER, ROOT, yaml_load
+from ultralytics.yolo.utils.checks import check_requirements, check_suffix, check_version, check_yaml
+from ultralytics.yolo.utils.downloads import attempt_download_asset, is_url
+from ultralytics.yolo.utils.ops import xywh2xyxy
+
+
+def check_class_names(names):
+    """Check class names. Map imagenet class codes to human-readable names if required. Convert lists to dicts."""
+    if isinstance(names, list):  # names is a list
+        names = dict(enumerate(names))  # convert to dict
+    if isinstance(names, dict):
+        # Convert 1) string keys to int, i.e. '0' to 0, and non-string values to strings, i.e. True to 'True'
+        names = {int(k): str(v) for k, v in names.items()}
+        n = len(names)
+        if max(names.keys()) >= n:
+            raise KeyError(f'{n}-class dataset requires class indices 0-{n - 1}, but you have invalid class indices '
+                           f'{min(names.keys())}-{max(names.keys())} defined in your dataset YAML.')
+        if isinstance(names[0], str) and names[0].startswith('n0'):  # imagenet class codes, i.e. 'n01440764'
+            map = yaml_load(ROOT / 'datasets/ImageNet.yaml')['map']  # human-readable names
+            names = {k: map[v] for k, v in names.items()}
+    return names
+
+
+class AutoBackend(nn.Module):
+
+    def __init__(self,
+                 weights='yolov8n.pt',
+                 device=torch.device('cpu'),
+                 dnn=False,
+                 data=None,
+                 fp16=False,
+                 fuse=True,
+                 verbose=True):
+        """
+        MultiBackend class for python inference on various platforms using Ultralytics YOLO.
+
+        Args:
+            weights (str): The path to the weights file. Default: 'yolov8n.pt'
+            device (torch.device): The device to run the model on.
+            dnn (bool): Use OpenCV DNN module for inference if True, defaults to False.
+            data (str | Path | optional): Additional data.yaml file for class names.
+            fp16 (bool): If True, use half precision. Default: False
+            fuse (bool): Whether to fuse the model or not. Default: True
+            verbose (bool): Whether to run in verbose mode or not. Default: True
+
+        Supported formats and their naming conventions:
+            | Format                | Suffix           |
+            |-----------------------|------------------|
+            | PyTorch               | *.pt             |
+            | TorchScript           | *.torchscript    |
+            | ONNX Runtime          | *.onnx           |
+            | ONNX OpenCV DNN       | *.onnx dnn=True  |
+            | OpenVINO              | *.xml            |
+            | CoreML                | *.mlmodel        |
+            | TensorRT              | *.engine         |
+            | TensorFlow SavedModel | *_saved_model    |
+            | TensorFlow GraphDef   | *.pb             |
+            | TensorFlow Lite       | *.tflite         |
+            | TensorFlow Edge TPU   | *_edgetpu.tflite |
+            | PaddlePaddle          | *_paddle_model   |
+        """
+        super().__init__()
+        w = str(weights[0] if isinstance(weights, list) else weights)
+        nn_module = isinstance(weights, torch.nn.Module)
+        pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, triton = self._model_type(w)
+        fp16 &= pt or jit or onnx or engine or nn_module or triton  # FP16
+        nhwc = coreml or saved_model or pb or tflite or edgetpu  # BHWC formats (vs torch BCWH)
+        stride = 32  # default stride
+        model, metadata = None, None
+        cuda = torch.cuda.is_available() and device.type != 'cpu'  # use CUDA
+        if not (pt or triton or nn_module):
+            w = attempt_download_asset(w)  # download if not local
+
+        # NOTE: special case: in-memory pytorch model
+        if nn_module:
+            model = weights.to(device)
+            model = model.fuse(verbose=verbose) if fuse else model
+            if hasattr(model, 'kpt_shape'):
+                kpt_shape = model.kpt_shape  # pose-only
+            stride = max(int(model.stride.max()), 32)  # model stride
+            names = model.module.names if hasattr(model, 'module') else model.names  # get class names
+            model.half() if fp16 else model.float()
+            self.model = model  # explicitly assign for to(), cpu(), cuda(), half()
+            pt = True
+        elif pt:  # PyTorch
+            from ultralytics.nn.tasks import attempt_load_weights
+            model = attempt_load_weights(weights if isinstance(weights, list) else w,
+                                         device=device,
+                                         inplace=True,
+                                         fuse=fuse)
+            if hasattr(model, 'kpt_shape'):
+                kpt_shape = model.kpt_shape  # pose-only
+            stride = max(int(model.stride.max()), 32)  # model stride
+            names = model.module.names if hasattr(model, 'module') else model.names  # get class names
+            model.half() if fp16 else model.float()
+            self.model = model  # explicitly assign for to(), cpu(), cuda(), half()
+        elif jit:  # TorchScript
+            LOGGER.info(f'Loading {w} for TorchScript inference...')
+            extra_files = {'config.txt': ''}  # model metadata
+            model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
+            model.half() if fp16 else model.float()
+            if extra_files['config.txt']:  # load metadata dict
+                metadata = json.loads(extra_files['config.txt'], object_hook=lambda x: dict(x.items()))
+        elif dnn:  # ONNX OpenCV DNN
+            LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
+            check_requirements('opencv-python>=4.5.4')
+            net = cv2.dnn.readNetFromONNX(w)
+        elif onnx:  # ONNX Runtime
+            LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
+            check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
+            import onnxruntime
+            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
+            session = onnxruntime.InferenceSession(w, providers=providers)
+            output_names = [x.name for x in session.get_outputs()]
+            metadata = session.get_modelmeta().custom_metadata_map  # metadata
+        elif xml:  # OpenVINO
+            LOGGER.info(f'Loading {w} for OpenVINO inference...')
+            check_requirements('openvino')  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+            from openvino.runtime import Core, Layout, get_batch  # noqa
+            ie = Core()
+            w = Path(w)
+            if not w.is_file():  # if not *.xml
+                w = next(w.glob('*.xml'))  # get *.xml file from *_openvino_model dir
+            network = ie.read_model(model=str(w), weights=w.with_suffix('.bin'))
+            if network.get_parameters()[0].get_layout().empty:
+                network.get_parameters()[0].set_layout(Layout('NCHW'))
+            batch_dim = get_batch(network)
+            if batch_dim.is_static:
+                batch_size = batch_dim.get_length()
+            executable_network = ie.compile_model(network, device_name='CPU')  # device_name="MYRIAD" for NCS2
+            metadata = w.parent / 'metadata.yaml'
+        elif engine:  # TensorRT
+            LOGGER.info(f'Loading {w} for TensorRT inference...')
+            try:
+                import tensorrt as trt  # noqa https://developer.nvidia.com/nvidia-tensorrt-download
+            except ImportError:
+                if LINUX:
+                    check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com')
+                import tensorrt as trt  # noqa
+            check_version(trt.__version__, '7.0.0', hard=True)  # require tensorrt>=7.0.0
+            if device.type == 'cpu':
+                device = torch.device('cuda:0')
+            Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
+            logger = trt.Logger(trt.Logger.INFO)
+            # Read file
+            with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
+                meta_len = int.from_bytes(f.read(4), byteorder='little')  # read metadata length
+                metadata = json.loads(f.read(meta_len).decode('utf-8'))  # read metadata
+                model = runtime.deserialize_cuda_engine(f.read())  # read engine
+            context = model.create_execution_context()
+            bindings = OrderedDict()
+            output_names = []
+            fp16 = False  # default updated below
+            dynamic = False
+            for i in range(model.num_bindings):
+                name = model.get_binding_name(i)
+                dtype = trt.nptype(model.get_binding_dtype(i))
+                if model.binding_is_input(i):
+                    if -1 in tuple(model.get_binding_shape(i)):  # dynamic
+                        dynamic = True
+                        context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[2]))
+                    if dtype == np.float16:
+                        fp16 = True
+                else:  # output
+                    output_names.append(name)
+                shape = tuple(context.get_binding_shape(i))
+                im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
+                bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
+            binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
+            batch_size = bindings['images'].shape[0]  # if dynamic, this is instead max batch size
+        elif coreml:  # CoreML
+            LOGGER.info(f'Loading {w} for CoreML inference...')
+            import coremltools as ct
+            model = ct.models.MLModel(w)
+            metadata = dict(model.user_defined_metadata)
+        elif saved_model:  # TF SavedModel
+            LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...')
+            import tensorflow as tf
+            keras = False  # assume TF1 saved_model
+            model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
+            metadata = Path(w) / 'metadata.yaml'
+        elif pb:  # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
+            LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')
+            import tensorflow as tf
+
+            from ultralytics.yolo.engine.exporter import gd_outputs
+
+            def wrap_frozen_graph(gd, inputs, outputs):
+                """Wrap frozen graphs for deployment."""
+                x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=''), [])  # wrapped
+                ge = x.graph.as_graph_element
+                return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
+
+            gd = tf.Graph().as_graph_def()  # TF GraphDef
+            with open(w, 'rb') as f:
+                gd.ParseFromString(f.read())
+            frozen_func = wrap_frozen_graph(gd, inputs='x:0', outputs=gd_outputs(gd))
+        elif tflite or edgetpu:  # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
+            try:  # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
+                from tflite_runtime.interpreter import Interpreter, load_delegate
+            except ImportError:
+                import tensorflow as tf
+                Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate
+            if edgetpu:  # TF Edge TPU https://coral.ai/software/#edgetpu-runtime
+                LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
+                delegate = {
+                    'Linux': 'libedgetpu.so.1',
+                    'Darwin': 'libedgetpu.1.dylib',
+                    'Windows': 'edgetpu.dll'}[platform.system()]
+                interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
+            else:  # TFLite
+                LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
+                interpreter = Interpreter(model_path=w)  # load TFLite model
+            interpreter.allocate_tensors()  # allocate
+            input_details = interpreter.get_input_details()  # inputs
+            output_details = interpreter.get_output_details()  # outputs
+            # Load metadata
+            with contextlib.suppress(zipfile.BadZipFile):
+                with zipfile.ZipFile(w, 'r') as model:
+                    meta_file = model.namelist()[0]
+                    metadata = ast.literal_eval(model.read(meta_file).decode('utf-8'))
+        elif tfjs:  # TF.js
+            raise NotImplementedError('YOLOv8 TF.js inference is not supported')
+        elif paddle:  # PaddlePaddle
+            LOGGER.info(f'Loading {w} for PaddlePaddle inference...')
+            check_requirements('paddlepaddle-gpu' if cuda else 'paddlepaddle')
+            import paddle.inference as pdi  # noqa
+            w = Path(w)
+            if not w.is_file():  # if not *.pdmodel
+                w = next(w.rglob('*.pdmodel'))  # get *.pdmodel file from *_paddle_model dir
+            config = pdi.Config(str(w), str(w.with_suffix('.pdiparams')))
+            if cuda:
+                config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0)
+            predictor = pdi.create_predictor(config)
+            input_handle = predictor.get_input_handle(predictor.get_input_names()[0])
+            output_names = predictor.get_output_names()
+            metadata = w.parents[1] / 'metadata.yaml'
+        elif triton:  # NVIDIA Triton Inference Server
+            LOGGER.info('Triton Inference Server not supported...')
+            '''
+            TODO:
+            check_requirements('tritonclient[all]')
+            from utils.triton import TritonRemoteModel
+            model = TritonRemoteModel(url=w)
+            nhwc = model.runtime.startswith("tensorflow")
+            '''
+        else:
+            from ultralytics.yolo.engine.exporter import export_formats
+            raise TypeError(f"model='{w}' is not a supported model format. "
+                            'See https://docs.ultralytics.com/modes/predict for help.'
+                            f'\n\n{export_formats()}')
+
+        # Load external metadata YAML
+        if isinstance(metadata, (str, Path)) and Path(metadata).exists():
+            metadata = yaml_load(metadata)
+        if metadata:
+            for k, v in metadata.items():
+                if k in ('stride', 'batch'):
+                    metadata[k] = int(v)
+                elif k in ('imgsz', 'names', 'kpt_shape') and isinstance(v, str):
+                    metadata[k] = eval(v)
+            stride = metadata['stride']
+            task = metadata['task']
+            batch = metadata['batch']
+            imgsz = metadata['imgsz']
+            names = metadata['names']
+            kpt_shape = metadata.get('kpt_shape')
+        elif not (pt or triton or nn_module):
+            LOGGER.warning(f"WARNING ⚠️ Metadata not found for 'model={weights}'")
+
+        # Check names
+        if 'names' not in locals():  # names missing
+            names = self._apply_default_class_names(data)
+        names = check_class_names(names)
+
+        self.__dict__.update(locals())  # assign all variables to self
+
+    def forward(self, im, augment=False, visualize=False):
+        """
+        Runs inference on the YOLOv8 MultiBackend model.
+
+        Args:
+            im (torch.Tensor): The image tensor to perform inference on.
+            augment (bool): whether to perform data augmentation during inference, defaults to False
+            visualize (bool): whether to visualize the output predictions, defaults to False
+
+        Returns:
+            (tuple): Tuple containing the raw output tensor, and processed output for visualization (if visualize=True)
+        """
+        b, ch, h, w = im.shape  # batch, channel, height, width
+        if self.fp16 and im.dtype != torch.float16:
+            im = im.half()  # to FP16
+        if self.nhwc:
+            im = im.permute(0, 2, 3, 1)  # torch BCHW to numpy BHWC shape(1,320,192,3)
+
+        if self.pt or self.nn_module:  # PyTorch
+            y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)
+        elif self.jit:  # TorchScript
+            y = self.model(im)
+        elif self.dnn:  # ONNX OpenCV DNN
+            im = im.cpu().numpy()  # torch to numpy
+            self.net.setInput(im)
+            y = self.net.forward()
+        elif self.onnx:  # ONNX Runtime
+            im = im.cpu().numpy()  # torch to numpy
+            y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
+        elif self.xml:  # OpenVINO
+            im = im.cpu().numpy()  # FP32
+            y = list(self.executable_network([im]).values())
+        elif self.engine:  # TensorRT
+            if self.dynamic and im.shape != self.bindings['images'].shape:
+                i = self.model.get_binding_index('images')
+                self.context.set_binding_shape(i, im.shape)  # reshape if dynamic
+                self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)
+                for name in self.output_names:
+                    i = self.model.get_binding_index(name)
+                    self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))
+            s = self.bindings['images'].shape
+            assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
+            self.binding_addrs['images'] = int(im.data_ptr())
+            self.context.execute_v2(list(self.binding_addrs.values()))
+            y = [self.bindings[x].data for x in sorted(self.output_names)]
+        elif self.coreml:  # CoreML
+            im = im[0].cpu().numpy()
+            im_pil = Image.fromarray((im * 255).astype('uint8'))
+            # im = im.resize((192, 320), Image.ANTIALIAS)
+            y = self.model.predict({'image': im_pil})  # coordinates are xywh normalized
+            if 'confidence' in y:
+                box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]])  # xyxy pixels
+                conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
+                y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
+            elif len(y) == 1:  # classification model
+                y = list(y.values())
+            elif len(y) == 2:  # segmentation model
+                y = list(reversed(y.values()))  # reversed for segmentation models (pred, proto)
+        elif self.paddle:  # PaddlePaddle
+            im = im.cpu().numpy().astype(np.float32)
+            self.input_handle.copy_from_cpu(im)
+            self.predictor.run()
+            y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names]
+        elif self.triton:  # NVIDIA Triton Inference Server
+            y = self.model(im)
+        else:  # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
+            im = im.cpu().numpy()
+            if self.saved_model:  # SavedModel
+                y = self.model(im, training=False) if self.keras else self.model(im)
+                if not isinstance(y, list):
+                    y = [y]
+            elif self.pb:  # GraphDef
+                y = self.frozen_func(x=self.tf.constant(im))
+                if len(y) == 2 and len(self.names) == 999:  # segments and names not defined
+                    ip, ib = (0, 1) if len(y[0].shape) == 4 else (1, 0)  # index of protos, boxes
+                    nc = y[ib].shape[1] - y[ip].shape[3] - 4  # y = (1, 160, 160, 32), (1, 116, 8400)
+                    self.names = {i: f'class{i}' for i in range(nc)}
+            else:  # Lite or Edge TPU
+                input = self.input_details[0]
+                int8 = input['dtype'] == np.int8  # is TFLite quantized int8 model
+                if int8:
+                    scale, zero_point = input['quantization']
+                    im = (im / scale + zero_point).astype(np.int8)  # de-scale
+                self.interpreter.set_tensor(input['index'], im)
+                self.interpreter.invoke()
+                y = []
+                for output in self.output_details:
+                    x = self.interpreter.get_tensor(output['index'])
+                    if int8:
+                        scale, zero_point = output['quantization']
+                        x = (x.astype(np.float32) - zero_point) * scale  # re-scale
+                    y.append(x)
+            # TF segment fixes: export is reversed vs ONNX export and protos are transposed
+            if len(y) == 2:  # segment with (det, proto) output order reversed
+                if len(y[1].shape) != 4:
+                    y = list(reversed(y))  # should be y = (1, 116, 8400), (1, 160, 160, 32)
+                y[1] = np.transpose(y[1], (0, 3, 1, 2))  # should be y = (1, 116, 8400), (1, 32, 160, 160)
+            y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
+            # y[0][..., :4] *= [w, h, w, h]  # xywh normalized to pixels
+
+        # for x in y:
+        #     print(type(x), len(x)) if isinstance(x, (list, tuple)) else print(type(x), x.shape)  # debug shapes
+        if isinstance(y, (list, tuple)):
+            return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
+        else:
+            return self.from_numpy(y)
+
+    def from_numpy(self, x):
+        """
+         Convert a numpy array to a tensor.
+
+         Args:
+             x (np.ndarray): The array to be converted.
+
+         Returns:
+             (torch.Tensor): The converted tensor
+         """
+        return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
+
+    def warmup(self, imgsz=(1, 3, 640, 640)):
+        """
+        Warm up the model by running one forward pass with a dummy input.
+
+        Args:
+            imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width)
+
+        Returns:
+            (None): This method runs the forward pass and don't return any value
+        """
+        warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module
+        if any(warmup_types) and (self.device.type != 'cpu' or self.triton):
+            im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device)  # input
+            for _ in range(2 if self.jit else 1):  #
+                self.forward(im)  # warmup
+
+    @staticmethod
+    def _apply_default_class_names(data):
+        """Applies default class names to an input YAML file or returns numerical class names."""
+        with contextlib.suppress(Exception):
+            return yaml_load(check_yaml(data))['names']
+        return {i: f'class{i}' for i in range(999)}  # return default if above errors
+
+    @staticmethod
+    def _model_type(p='path/to/model.pt'):
+        """
+        This function takes a path to a model file and returns the model type
+
+        Args:
+            p: path to the model file. Defaults to path/to/model.pt
+        """
+        # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
+        # types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]
+        from ultralytics.yolo.engine.exporter import export_formats
+        sf = list(export_formats().Suffix)  # export suffixes
+        if not is_url(p, check=False) and not isinstance(p, str):
+            check_suffix(p, sf)  # checks
+        url = urlparse(p)  # if url may be Triton inference server
+        types = [s in Path(p).name for s in sf]
+        types[8] &= not types[9]  # tflite &= not edgetpu
+        triton = not any(types) and all([any(s in url.scheme for s in ['http', 'grpc']), url.netloc])
+        return types + [triton]

ultralytics/nn/autoshape.py

@@ -0,0 +1,244 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Common modules
+"""
+
+from copy import copy
+from pathlib import Path
+
+import cv2
+import numpy as np
+import requests
+import torch
+import torch.nn as nn
+from PIL import Image, ImageOps
+from torch.cuda import amp
+
+from ultralytics.nn.autobackend import AutoBackend
+from ultralytics.yolo.data.augment import LetterBox
+from ultralytics.yolo.utils import LOGGER, colorstr
+from ultralytics.yolo.utils.files import increment_path
+from ultralytics.yolo.utils.ops import Profile, make_divisible, non_max_suppression, scale_boxes, xyxy2xywh
+from ultralytics.yolo.utils.plotting import Annotator, colors, save_one_box
+from ultralytics.yolo.utils.torch_utils import copy_attr, smart_inference_mode
+
+
+class AutoShape(nn.Module):
+    """YOLOv8 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
+    agnostic = False  # NMS class-agnostic
+    multi_label = False  # NMS multiple labels per box
+    classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
+    max_det = 1000  # maximum number of detections per image
+    amp = False  # Automatic Mixed Precision (AMP) inference
+
+    def __init__(self, model, verbose=True):
+        """Initializes object and copies attributes from model object."""
+        super().__init__()
+        if verbose:
+            LOGGER.info('Adding AutoShape... ')
+        copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=())  # copy attributes
+        self.dmb = isinstance(model, AutoBackend)  # DetectMultiBackend() instance
+        self.pt = not self.dmb or model.pt  # PyTorch model
+        self.model = model.eval()
+        if self.pt:
+            m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
+            m.inplace = False  # Detect.inplace=False for safe multithread inference
+            m.export = True  # do not output loss values
+
+    def _apply(self, fn):
+        """Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers."""
+        self = super()._apply(fn)
+        if self.pt:
+            m = self.model.model.model[-1] if self.dmb else self.model.model[-1]  # Detect()
+            m.stride = fn(m.stride)
+            m.grid = list(map(fn, m.grid))
+            if isinstance(m.anchor_grid, list):
+                m.anchor_grid = list(map(fn, m.anchor_grid))
+        return self
+
+    @smart_inference_mode()
+    def forward(self, ims, size=640, augment=False, profile=False):
+        """Inference from various sources. For size(height=640, width=1280), RGB images example inputs are:."""
+        #   file:        ims = 'data/images/zidane.jpg'  # str or PosixPath
+        #   URI:             = 'https://ultralytics.com/images/zidane.jpg'
+        #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(640,1280,3)
+        #   PIL:             = Image.open('image.jpg') or ImageGrab.grab()  # HWC x(640,1280,3)
+        #   numpy:           = np.zeros((640,1280,3))  # HWC
+        #   torch:           = torch.zeros(16,3,320,640)  # BCHW (scaled to size=640, 0-1 values)
+        #   multiple:        = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
+
+        dt = (Profile(), Profile(), Profile())
+        with dt[0]:
+            if isinstance(size, int):  # expand
+                size = (size, size)
+            p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device)  # param
+            autocast = self.amp and (p.device.type != 'cpu')  # Automatic Mixed Precision (AMP) inference
+            if isinstance(ims, torch.Tensor):  # torch
+                with amp.autocast(autocast):
+                    return self.model(ims.to(p.device).type_as(p), augment=augment)  # inference
+
+            # Preprocess
+            n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims])  # number, list of images
+            shape0, shape1, files = [], [], []  # image and inference shapes, filenames
+            for i, im in enumerate(ims):
+                f = f'image{i}'  # filename
+                if isinstance(im, (str, Path)):  # filename or uri
+                    im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
+                    im = np.asarray(ImageOps.exif_transpose(im))
+                elif isinstance(im, Image.Image):  # PIL Image
+                    im, f = np.asarray(ImageOps.exif_transpose(im)), getattr(im, 'filename', f) or f
+                files.append(Path(f).with_suffix('.jpg').name)
+                if im.shape[0] < 5:  # image in CHW
+                    im = im.transpose((1, 2, 0))  # reverse dataloader .transpose(2, 0, 1)
+                im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)  # enforce 3ch input
+                s = im.shape[:2]  # HWC
+                shape0.append(s)  # image shape
+                g = max(size) / max(s)  # gain
+                shape1.append([y * g for y in s])
+                ims[i] = im if im.data.contiguous else np.ascontiguousarray(im)  # update
+            shape1 = [make_divisible(x, self.stride) for x in np.array(shape1).max(0)] if self.pt else size  # inf shape
+            x = [LetterBox(shape1, auto=False)(image=im)['img'] for im in ims]  # pad
+            x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2)))  # stack and BHWC to BCHW
+            x = torch.from_numpy(x).to(p.device).type_as(p) / 255  # uint8 to fp16/32
+
+        with amp.autocast(autocast):
+            # Inference
+            with dt[1]:
+                y = self.model(x, augment=augment)  # forward
+
+            # Postprocess
+            with dt[2]:
+                y = non_max_suppression(y if self.dmb else y[0],
+                                        self.conf,
+                                        self.iou,
+                                        self.classes,
+                                        self.agnostic,
+                                        self.multi_label,
+                                        max_det=self.max_det)  # NMS
+                for i in range(n):
+                    scale_boxes(shape1, y[i][:, :4], shape0[i])
+
+            return Detections(ims, y, files, dt, self.names, x.shape)
+
+
+class Detections:
+    """ YOLOv8 detections class for inference results"""
+
+    def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None):
+        """Initialize object attributes for YOLO detection results."""
+        super().__init__()
+        d = pred[0].device  # device
+        gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in ims]  # normalizations
+        self.ims = ims  # list of images as numpy arrays
+        self.pred = pred  # list of tensors pred[0] = (xyxy, conf, cls)
+        self.names = names  # class names
+        self.files = files  # image filenames
+        self.times = times  # profiling times
+        self.xyxy = pred  # xyxy pixels
+        self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
+        self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
+        self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
+        self.n = len(self.pred)  # number of images (batch size)
+        self.t = tuple(x.t / self.n * 1E3 for x in times)  # timestamps (ms)
+        self.s = tuple(shape)  # inference BCHW shape
+
+    def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):
+        """Return performance metrics and optionally cropped/save images or results."""
+        s, crops = '', []
+        for i, (im, pred) in enumerate(zip(self.ims, self.pred)):
+            s += f'\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '  # string
+            if pred.shape[0]:
+                for c in pred[:, -1].unique():
+                    n = (pred[:, -1] == c).sum()  # detections per class
+                    s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                s = s.rstrip(', ')
+                if show or save or render or crop:
+                    annotator = Annotator(im, example=str(self.names))
+                    for *box, conf, cls in reversed(pred):  # xyxy, confidence, class
+                        label = f'{self.names[int(cls)]} {conf:.2f}'
+                        if crop:
+                            file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
+                            crops.append({
+                                'box': box,
+                                'conf': conf,
+                                'cls': cls,
+                                'label': label,
+                                'im': save_one_box(box, im, file=file, save=save)})
+                        else:  # all others
+                            annotator.box_label(box, label if labels else '', color=colors(cls))
+                    im = annotator.im
+            else:
+                s += '(no detections)'
+
+            im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im  # from np
+            if show:
+                im.show(self.files[i])  # show
+            if save:
+                f = self.files[i]
+                im.save(save_dir / f)  # save
+                if i == self.n - 1:
+                    LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
+            if render:
+                self.ims[i] = np.asarray(im)
+        if pprint:
+            s = s.lstrip('\n')
+            return f'{s}\nSpeed: %.1fms preprocess, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t
+        if crop:
+            if save:
+                LOGGER.info(f'Saved results to {save_dir}\n')
+            return crops
+
+    def show(self, labels=True):
+        """Displays YOLO results with detected bounding boxes."""
+        self._run(show=True, labels=labels)  # show results
+
+    def save(self, labels=True, save_dir='runs/detect/exp', exist_ok=False):
+        """Save detection results with optional labels to specified directory."""
+        save_dir = increment_path(save_dir, exist_ok, mkdir=True)  # increment save_dir
+        self._run(save=True, labels=labels, save_dir=save_dir)  # save results
+
+    def crop(self, save=True, save_dir='runs/detect/exp', exist_ok=False):
+        """Crops images into detections and saves them if 'save' is True."""
+        save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None
+        return self._run(crop=True, save=save, save_dir=save_dir)  # crop results
+
+    def render(self, labels=True):
+        """Renders detected objects and returns images."""
+        self._run(render=True, labels=labels)  # render results
+        return self.ims
+
+    def pandas(self):
+        """Return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])."""
+        import pandas
+        new = copy(self)  # return copy
+        ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name'  # xyxy columns
+        cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name'  # xywh columns
+        for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
+            a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)]  # update
+            setattr(new, k, [pandas.DataFrame(x, columns=c) for x in a])
+        return new
+
+    def tolist(self):
+        """Return a list of Detections objects, i.e. 'for result in results.tolist():'."""
+        r = range(self.n)  # iterable
+        x = [Detections([self.ims[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]
+        # for d in x:
+        #    for k in ['ims', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
+        #        setattr(d, k, getattr(d, k)[0])  # pop out of list
+        return x
+
+    def print(self):
+        """Print the results of the `self._run()` function."""
+        LOGGER.info(self.__str__())
+
+    def __len__(self):  # override len(results)
+        return self.n
+
+    def __str__(self):  # override print(results)
+        return self._run(pprint=True)  # print results
+
+    def __repr__(self):
+        """Returns a printable representation of the object."""
+        return f'YOLOv8 {self.__class__} instance\n' + self.__str__()

ultralytics/nn/modules/__init__.py

@@ -0,0 +1,29 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Ultralytics modules. Visualize with:
+
+from ultralytics.nn.modules import *
+import torch
+import os
+
+x = torch.ones(1, 128, 40, 40)
+m = Conv(128, 128)
+f = f'{m._get_name()}.onnx'
+torch.onnx.export(m, x, f)
+os.system(f'onnxsim {f} {f} && open {f}')
+"""
+
+from .block import (C1, C2, C3, C3TR, DFL, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x, GhostBottleneck,
+                    HGBlock, HGStem, Proto, RepC3)
+from .conv import (CBAM, ChannelAttention, Concat, Conv, Conv2, ConvTranspose, DWConv, DWConvTranspose2d, Focus,
+                   GhostConv, LightConv, RepConv, SpatialAttention)
+from .head import Classify, Detect, Pose, RTDETRDecoder, Segment
+from .transformer import (AIFI, MLP, DeformableTransformerDecoder, DeformableTransformerDecoderLayer, LayerNorm2d,
+                          MLPBlock, MSDeformAttn, TransformerBlock, TransformerEncoderLayer, TransformerLayer)
+
+__all__ = ('Conv', 'Conv2', 'LightConv', 'RepConv', 'DWConv', 'DWConvTranspose2d', 'ConvTranspose', 'Focus',
+           'GhostConv', 'ChannelAttention', 'SpatialAttention', 'CBAM', 'Concat', 'TransformerLayer',
+           'TransformerBlock', 'MLPBlock', 'LayerNorm2d', 'DFL', 'HGBlock', 'HGStem', 'SPP', 'SPPF', 'C1', 'C2', 'C3',
+           'C2f', 'C3x', 'C3TR', 'C3Ghost', 'GhostBottleneck', 'Bottleneck', 'BottleneckCSP', 'Proto', 'Detect',
+           'Segment', 'Pose', 'Classify', 'TransformerEncoderLayer', 'RepC3', 'RTDETRDecoder', 'AIFI',
+           'DeformableTransformerDecoder', 'DeformableTransformerDecoderLayer', 'MSDeformAttn', 'MLP')

ultralytics/nn/modules/block.py

@@ -0,0 +1,304 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Block modules
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .conv import Conv, DWConv, GhostConv, LightConv, RepConv
+from .transformer import TransformerBlock
+
+__all__ = ('DFL', 'HGBlock', 'HGStem', 'SPP', 'SPPF', 'C1', 'C2', 'C3', 'C2f', 'C3x', 'C3TR', 'C3Ghost',
+           'GhostBottleneck', 'Bottleneck', 'BottleneckCSP', 'Proto', 'RepC3')
+
+
+class DFL(nn.Module):
+    """
+    Integral module of Distribution Focal Loss (DFL).
+    Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
+    """
+
+    def __init__(self, c1=16):
+        """Initialize a convolutional layer with a given number of input channels."""
+        super().__init__()
+        self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
+        x = torch.arange(c1, dtype=torch.float)
+        self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
+        self.c1 = c1
+
+    def forward(self, x):
+        """Applies a transformer layer on input tensor 'x' and returns a tensor."""
+        b, c, a = x.shape  # batch, channels, anchors
+        return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
+        # return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
+
+
+class Proto(nn.Module):
+    """YOLOv8 mask Proto module for segmentation models."""
+
+    def __init__(self, c1, c_=256, c2=32):  # ch_in, number of protos, number of masks
+        super().__init__()
+        self.cv1 = Conv(c1, c_, k=3)
+        self.upsample = nn.ConvTranspose2d(c_, c_, 2, 2, 0, bias=True)  # nn.Upsample(scale_factor=2, mode='nearest')
+        self.cv2 = Conv(c_, c_, k=3)
+        self.cv3 = Conv(c_, c2)
+
+    def forward(self, x):
+        """Performs a forward pass through layers using an upsampled input image."""
+        return self.cv3(self.cv2(self.upsample(self.cv1(x))))
+
+
+class HGStem(nn.Module):
+    """StemBlock of PPHGNetV2 with 5 convolutions and one maxpool2d.
+    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
+    """
+
+    def __init__(self, c1, cm, c2):
+        super().__init__()
+        self.stem1 = Conv(c1, cm, 3, 2, act=nn.ReLU())
+        self.stem2a = Conv(cm, cm // 2, 2, 1, 0, act=nn.ReLU())
+        self.stem2b = Conv(cm // 2, cm, 2, 1, 0, act=nn.ReLU())
+        self.stem3 = Conv(cm * 2, cm, 3, 2, act=nn.ReLU())
+        self.stem4 = Conv(cm, c2, 1, 1, act=nn.ReLU())
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=1, padding=0, ceil_mode=True)
+
+    def forward(self, x):
+        """Forward pass of a PPHGNetV2 backbone layer."""
+        x = self.stem1(x)
+        x = F.pad(x, [0, 1, 0, 1])
+        x2 = self.stem2a(x)
+        x2 = F.pad(x2, [0, 1, 0, 1])
+        x2 = self.stem2b(x2)
+        x1 = self.pool(x)
+        x = torch.cat([x1, x2], dim=1)
+        x = self.stem3(x)
+        x = self.stem4(x)
+        return x
+
+
+class HGBlock(nn.Module):
+    """HG_Block of PPHGNetV2 with 2 convolutions and LightConv.
+    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
+    """
+
+    def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
+        super().__init__()
+        block = LightConv if lightconv else Conv
+        self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
+        self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act)  # squeeze conv
+        self.ec = Conv(c2 // 2, c2, 1, 1, act=act)  # excitation conv
+        self.add = shortcut and c1 == c2
+
+    def forward(self, x):
+        """Forward pass of a PPHGNetV2 backbone layer."""
+        y = [x]
+        y.extend(m(y[-1]) for m in self.m)
+        y = self.ec(self.sc(torch.cat(y, 1)))
+        return y + x if self.add else y
+
+
+class SPP(nn.Module):
+    """Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729."""
+
+    def __init__(self, c1, c2, k=(5, 9, 13)):
+        """Initialize the SPP layer with input/output channels and pooling kernel sizes."""
+        super().__init__()
+        c_ = c1 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
+        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
+
+    def forward(self, x):
+        """Forward pass of the SPP layer, performing spatial pyramid pooling."""
+        x = self.cv1(x)
+        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
+
+
+class SPPF(nn.Module):
+    """Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher."""
+
+    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
+        super().__init__()
+        c_ = c1 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c_ * 4, c2, 1, 1)
+        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
+
+    def forward(self, x):
+        """Forward pass through Ghost Convolution block."""
+        x = self.cv1(x)
+        y1 = self.m(x)
+        y2 = self.m(y1)
+        return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
+
+
+class C1(nn.Module):
+    """CSP Bottleneck with 1 convolution."""
+
+    def __init__(self, c1, c2, n=1):  # ch_in, ch_out, number
+        super().__init__()
+        self.cv1 = Conv(c1, c2, 1, 1)
+        self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))
+
+    def forward(self, x):
+        """Applies cross-convolutions to input in the C3 module."""
+        y = self.cv1(x)
+        return self.m(y) + y
+
+
+class C2(nn.Module):
+    """CSP Bottleneck with 2 convolutions."""
+
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        self.c = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
+        self.cv2 = Conv(2 * self.c, c2, 1)  # optional act=FReLU(c2)
+        # self.attention = ChannelAttention(2 * self.c)  # or SpatialAttention()
+        self.m = nn.Sequential(*(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)))
+
+    def forward(self, x):
+        """Forward pass through the CSP bottleneck with 2 convolutions."""
+        a, b = self.cv1(x).chunk(2, 1)
+        return self.cv2(torch.cat((self.m(a), b), 1))
+
+
+class C2f(nn.Module):
+    """CSP Bottleneck with 2 convolutions."""
+
+    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        self.c = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
+        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
+        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
+
+    def forward(self, x):
+        """Forward pass through C2f layer."""
+        y = list(self.cv1(x).chunk(2, 1))
+        y.extend(m(y[-1]) for m in self.m)
+        return self.cv2(torch.cat(y, 1))
+
+    def forward_split(self, x):
+        """Forward pass using split() instead of chunk()."""
+        y = list(self.cv1(x).split((self.c, self.c), 1))
+        y.extend(m(y[-1]) for m in self.m)
+        return self.cv2(torch.cat(y, 1))
+
+
+class C3(nn.Module):
+    """CSP Bottleneck with 3 convolutions."""
+
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c1, c_, 1, 1)
+        self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
+        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))
+
+    def forward(self, x):
+        """Forward pass through the CSP bottleneck with 2 convolutions."""
+        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
+
+
+class C3x(C3):
+    """C3 module with cross-convolutions."""
+
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        """Initialize C3TR instance and set default parameters."""
+        super().__init__(c1, c2, n, shortcut, g, e)
+        self.c_ = int(c2 * e)
+        self.m = nn.Sequential(*(Bottleneck(self.c_, self.c_, shortcut, g, k=((1, 3), (3, 1)), e=1) for _ in range(n)))
+
+
+class RepC3(nn.Module):
+    """Rep C3."""
+
+    def __init__(self, c1, c2, n=3, e=1.0):
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c2, 1, 1)
+        self.cv2 = Conv(c1, c2, 1, 1)
+        self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
+        self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()
+
+    def forward(self, x):
+        """Forward pass of RT-DETR neck layer."""
+        return self.cv3(self.m(self.cv1(x)) + self.cv2(x))
+
+
+class C3TR(C3):
+    """C3 module with TransformerBlock()."""
+
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        """Initialize C3Ghost module with GhostBottleneck()."""
+        super().__init__(c1, c2, n, shortcut, g, e)
+        c_ = int(c2 * e)
+        self.m = TransformerBlock(c_, c_, 4, n)
+
+
+class C3Ghost(C3):
+    """C3 module with GhostBottleneck()."""
+
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        """Initialize 'SPP' module with various pooling sizes for spatial pyramid pooling."""
+        super().__init__(c1, c2, n, shortcut, g, e)
+        c_ = int(c2 * e)  # hidden channels
+        self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
+
+
+class GhostBottleneck(nn.Module):
+    """Ghost Bottleneck https://github.com/huawei-noah/ghostnet."""
+
+    def __init__(self, c1, c2, k=3, s=1):  # ch_in, ch_out, kernel, stride
+        super().__init__()
+        c_ = c2 // 2
+        self.conv = nn.Sequential(
+            GhostConv(c1, c_, 1, 1),  # pw
+            DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),  # dw
+            GhostConv(c_, c2, 1, 1, act=False))  # pw-linear
+        self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1,
+                                                                            act=False)) if s == 2 else nn.Identity()
+
+    def forward(self, x):
+        """Applies skip connection and concatenation to input tensor."""
+        return self.conv(x) + self.shortcut(x)
+
+
+class Bottleneck(nn.Module):
+    """Standard bottleneck."""
+
+    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):  # ch_in, ch_out, shortcut, groups, kernels, expand
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, k[0], 1)
+        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
+        self.add = shortcut and c1 == c2
+
+    def forward(self, x):
+        """'forward()' applies the YOLOv5 FPN to input data."""
+        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
+
+
+class BottleneckCSP(nn.Module):
+    """CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks."""
+
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
+        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
+        self.cv4 = Conv(2 * c_, c2, 1, 1)
+        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
+        self.act = nn.SiLU()
+        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
+
+    def forward(self, x):
+        """Applies a CSP bottleneck with 3 convolutions."""
+        y1 = self.cv3(self.m(self.cv1(x)))
+        y2 = self.cv2(x)
+        return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))

ultralytics/nn/modules/conv.py

@@ -0,0 +1,297 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Convolution modules
+"""
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+__all__ = ('Conv', 'LightConv', 'DWConv', 'DWConvTranspose2d', 'ConvTranspose', 'Focus', 'GhostConv',
+           'ChannelAttention', 'SpatialAttention', 'CBAM', 'Concat', 'RepConv')
+
+
+def autopad(k, p=None, d=1):  # kernel, padding, dilation
+    """Pad to 'same' shape outputs."""
+    if d > 1:
+        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p
+
+
+class Conv(nn.Module):
+    """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
+        """Initialize Conv layer with given arguments including activation."""
+        super().__init__()
+        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+    def forward(self, x):
+        """Apply convolution, batch normalization and activation to input tensor."""
+        return self.act(self.bn(self.conv(x)))
+
+    def forward_fuse(self, x):
+        """Perform transposed convolution of 2D data."""
+        return self.act(self.conv(x))
+
+
+class Conv2(Conv):
+    """Simplified RepConv module with Conv fusing."""
+
+    def __init__(self, c1, c2, k=3, s=1, p=None, g=1, d=1, act=True):
+        """Initialize Conv layer with given arguments including activation."""
+        super().__init__(c1, c2, k, s, p, g=g, d=d, act=act)
+        self.cv2 = nn.Conv2d(c1, c2, 1, s, autopad(1, p, d), groups=g, dilation=d, bias=False)  # add 1x1 conv
+
+    def forward(self, x):
+        """Apply convolution, batch normalization and activation to input tensor."""
+        return self.act(self.bn(self.conv(x) + self.cv2(x)))
+
+    def fuse_convs(self):
+        """Fuse parallel convolutions."""
+        w = torch.zeros_like(self.conv.weight.data)
+        i = [x // 2 for x in w.shape[2:]]
+        w[:, :, i[0]:i[0] + 1, i[1]:i[1] + 1] = self.cv2.weight.data.clone()
+        self.conv.weight.data += w
+        self.__delattr__('cv2')
+
+
+class LightConv(nn.Module):
+    """Light convolution with args(ch_in, ch_out, kernel).
+    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
+    """
+
+    def __init__(self, c1, c2, k=1, act=nn.ReLU()):
+        """Initialize Conv layer with given arguments including activation."""
+        super().__init__()
+        self.conv1 = Conv(c1, c2, 1, act=False)
+        self.conv2 = DWConv(c2, c2, k, act=act)
+
+    def forward(self, x):
+        """Apply 2 convolutions to input tensor."""
+        return self.conv2(self.conv1(x))
+
+
+class DWConv(Conv):
+    """Depth-wise convolution."""
+
+    def __init__(self, c1, c2, k=1, s=1, d=1, act=True):  # ch_in, ch_out, kernel, stride, dilation, activation
+        super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)
+
+
+class DWConvTranspose2d(nn.ConvTranspose2d):
+    """Depth-wise transpose convolution."""
+
+    def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0):  # ch_in, ch_out, kernel, stride, padding, padding_out
+        super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
+
+
+class ConvTranspose(nn.Module):
+    """Convolution transpose 2d layer."""
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, c1, c2, k=2, s=2, p=0, bn=True, act=True):
+        """Initialize ConvTranspose2d layer with batch normalization and activation function."""
+        super().__init__()
+        self.conv_transpose = nn.ConvTranspose2d(c1, c2, k, s, p, bias=not bn)
+        self.bn = nn.BatchNorm2d(c2) if bn else nn.Identity()
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+    def forward(self, x):
+        """Applies transposed convolutions, batch normalization and activation to input."""
+        return self.act(self.bn(self.conv_transpose(x)))
+
+    def forward_fuse(self, x):
+        """Applies activation and convolution transpose operation to input."""
+        return self.act(self.conv_transpose(x))
+
+
+class Focus(nn.Module):
+    """Focus wh information into c-space."""
+
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
+        super().__init__()
+        self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
+        # self.contract = Contract(gain=2)
+
+    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
+        return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
+        # return self.conv(self.contract(x))
+
+
+class GhostConv(nn.Module):
+    """Ghost Convolution https://github.com/huawei-noah/ghostnet."""
+
+    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):  # ch_in, ch_out, kernel, stride, groups
+        super().__init__()
+        c_ = c2 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
+        self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
+
+    def forward(self, x):
+        """Forward propagation through a Ghost Bottleneck layer with skip connection."""
+        y = self.cv1(x)
+        return torch.cat((y, self.cv2(y)), 1)
+
+
+class RepConv(nn.Module):
+    """RepConv is a basic rep-style block, including training and deploy status
+    This code is based on https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py
+    """
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, c1, c2, k=3, s=1, p=1, g=1, d=1, act=True, bn=False, deploy=False):
+        super().__init__()
+        assert k == 3 and p == 1
+        self.g = g
+        self.c1 = c1
+        self.c2 = c2
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+        self.bn = nn.BatchNorm2d(num_features=c1) if bn and c2 == c1 and s == 1 else None
+        self.conv1 = Conv(c1, c2, k, s, p=p, g=g, act=False)
+        self.conv2 = Conv(c1, c2, 1, s, p=(p - k // 2), g=g, act=False)
+
+    def forward_fuse(self, x):
+        """Forward process"""
+        return self.act(self.conv(x))
+
+    def forward(self, x):
+        """Forward process"""
+        id_out = 0 if self.bn is None else self.bn(x)
+        return self.act(self.conv1(x) + self.conv2(x) + id_out)
+
+    def get_equivalent_kernel_bias(self):
+        kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
+        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
+        kernelid, biasid = self._fuse_bn_tensor(self.bn)
+        return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
+
+    def _avg_to_3x3_tensor(self, avgp):
+        channels = self.c1
+        groups = self.g
+        kernel_size = avgp.kernel_size
+        input_dim = channels // groups
+        k = torch.zeros((channels, input_dim, kernel_size, kernel_size))
+        k[np.arange(channels), np.tile(np.arange(input_dim), groups), :, :] = 1.0 / kernel_size ** 2
+        return k
+
+    def _pad_1x1_to_3x3_tensor(self, kernel1x1):
+        if kernel1x1 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])
+
+    def _fuse_bn_tensor(self, branch):
+        if branch is None:
+            return 0, 0
+        if isinstance(branch, Conv):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        elif isinstance(branch, nn.BatchNorm2d):
+            if not hasattr(self, 'id_tensor'):
+                input_dim = self.c1 // self.g
+                kernel_value = np.zeros((self.c1, input_dim, 3, 3), dtype=np.float32)
+                for i in range(self.c1):
+                    kernel_value[i, i % input_dim, 1, 1] = 1
+                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def fuse_convs(self):
+        if hasattr(self, 'conv'):
+            return
+        kernel, bias = self.get_equivalent_kernel_bias()
+        self.conv = nn.Conv2d(in_channels=self.conv1.conv.in_channels,
+                              out_channels=self.conv1.conv.out_channels,
+                              kernel_size=self.conv1.conv.kernel_size,
+                              stride=self.conv1.conv.stride,
+                              padding=self.conv1.conv.padding,
+                              dilation=self.conv1.conv.dilation,
+                              groups=self.conv1.conv.groups,
+                              bias=True).requires_grad_(False)
+        self.conv.weight.data = kernel
+        self.conv.bias.data = bias
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__('conv1')
+        self.__delattr__('conv2')
+        if hasattr(self, 'nm'):
+            self.__delattr__('nm')
+        if hasattr(self, 'bn'):
+            self.__delattr__('bn')
+        if hasattr(self, 'id_tensor'):
+            self.__delattr__('id_tensor')
+
+
+class ChannelAttention(nn.Module):
+    """Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""
+
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
+        self.act = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.act(self.fc(self.pool(x)))
+
+
+class SpatialAttention(nn.Module):
+    """Spatial-attention module."""
+
+    def __init__(self, kernel_size=7):
+        """Initialize Spatial-attention module with kernel size argument."""
+        super().__init__()
+        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
+        padding = 3 if kernel_size == 7 else 1
+        self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.act = nn.Sigmoid()
+
+    def forward(self, x):
+        """Apply channel and spatial attention on input for feature recalibration."""
+        return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))
+
+
+class CBAM(nn.Module):
+    """Convolutional Block Attention Module."""
+
+    def __init__(self, c1, kernel_size=7):  # ch_in, kernels
+        super().__init__()
+        self.channel_attention = ChannelAttention(c1)
+        self.spatial_attention = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        """Applies the forward pass through C1 module."""
+        return self.spatial_attention(self.channel_attention(x))
+
+
+class Concat(nn.Module):
+    """Concatenate a list of tensors along dimension."""
+
+    def __init__(self, dimension=1):
+        """Concatenates a list of tensors along a specified dimension."""
+        super().__init__()
+        self.d = dimension
+
+    def forward(self, x):
+        """Forward pass for the YOLOv8 mask Proto module."""
+        return torch.cat(x, self.d)

ultralytics/nn/modules/head.py

@@ -0,0 +1,349 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Model head modules
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.init import constant_, xavier_uniform_
+
+from ultralytics.yolo.utils.tal import dist2bbox, make_anchors
+
+from .block import DFL, Proto
+from .conv import Conv
+from .transformer import MLP, DeformableTransformerDecoder, DeformableTransformerDecoderLayer
+from .utils import bias_init_with_prob, linear_init_
+
+__all__ = 'Detect', 'Segment', 'Pose', 'Classify', 'RTDETRDecoder'
+
+
+class Detect(nn.Module):
+    """YOLOv8 Detect head for detection models."""
+    dynamic = False  # force grid reconstruction
+    export = False  # export mode
+    shape = None
+    anchors = torch.empty(0)  # init
+    strides = torch.empty(0)  # init
+
+    def __init__(self, nc=80, ch=()):  # detection layer
+        super().__init__()
+        self.nc = nc  # number of classes
+        self.nl = len(ch)  # number of detection layers
+        self.reg_max = 16  # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
+        self.no = nc + self.reg_max * 4  # number of outputs per anchor
+        self.stride = torch.zeros(self.nl)  # strides computed during build
+        c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], self.nc)  # channels
+        self.cv2 = nn.ModuleList(
+            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
+        self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
+        self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
+
+    def forward(self, x):
+        """Concatenates and returns predicted bounding boxes and class probabilities."""
+        shape = x[0].shape  # BCHW
+        for i in range(self.nl):
+            x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
+        if self.training:
+            return x
+        elif self.dynamic or self.shape != shape:
+            self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
+            self.shape = shape
+
+        x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
+        if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'):  # avoid TF FlexSplitV ops
+            box = x_cat[:, :self.reg_max * 4]
+            cls = x_cat[:, self.reg_max * 4:]
+        else:
+            box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
+        dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
+        y = torch.cat((dbox, cls.sigmoid()), 1)
+        return y if self.export else (y, x)
+
+    def bias_init(self):
+        """Initialize Detect() biases, WARNING: requires stride availability."""
+        m = self  # self.model[-1]  # Detect() module
+        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
+        # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
+        for a, b, s in zip(m.cv2, m.cv3, m.stride):  # from
+            a[-1].bias.data[:] = 1.0  # box
+            b[-1].bias.data[:m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
+
+
+class Segment(Detect):
+    """YOLOv8 Segment head for segmentation models."""
+
+    def __init__(self, nc=80, nm=32, npr=256, ch=()):
+        """Initialize the YOLO model attributes such as the number of masks, prototypes, and the convolution layers."""
+        super().__init__(nc, ch)
+        self.nm = nm  # number of masks
+        self.npr = npr  # number of protos
+        self.proto = Proto(ch[0], self.npr, self.nm)  # protos
+        self.detect = Detect.forward
+
+        c4 = max(ch[0] // 4, self.nm)
+        self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nm, 1)) for x in ch)
+
+    def forward(self, x):
+        """Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
+        p = self.proto(x[0])  # mask protos
+        bs = p.shape[0]  # batch size
+
+        mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)  # mask coefficients
+        x = self.detect(self, x)
+        if self.training:
+            return x, mc, p
+        return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))
+
+
+class Pose(Detect):
+    """YOLOv8 Pose head for keypoints models."""
+
+    def __init__(self, nc=80, kpt_shape=(17, 3), ch=()):
+        """Initialize YOLO network with default parameters and Convolutional Layers."""
+        super().__init__(nc, ch)
+        self.kpt_shape = kpt_shape  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
+        self.nk = kpt_shape[0] * kpt_shape[1]  # number of keypoints total
+        self.detect = Detect.forward
+
+        c4 = max(ch[0] // 4, self.nk)
+        self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nk, 1)) for x in ch)
+
+    def forward(self, x):
+        """Perform forward pass through YOLO model and return predictions."""
+        bs = x[0].shape[0]  # batch size
+        kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1)  # (bs, 17*3, h*w)
+        x = self.detect(self, x)
+        if self.training:
+            return x, kpt
+        pred_kpt = self.kpts_decode(bs, kpt)
+        return torch.cat([x, pred_kpt], 1) if self.export else (torch.cat([x[0], pred_kpt], 1), (x[1], kpt))
+
+    def kpts_decode(self, bs, kpts):
+        """Decodes keypoints."""
+        ndim = self.kpt_shape[1]
+        if self.export:  # required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
+            y = kpts.view(bs, *self.kpt_shape, -1)
+            a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
+            if ndim == 3:
+                a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
+            return a.view(bs, self.nk, -1)
+        else:
+            y = kpts.clone()
+            if ndim == 3:
+                y[:, 2::3].sigmoid_()  # inplace sigmoid
+            y[:, 0::ndim] = (y[:, 0::ndim] * 2.0 + (self.anchors[0] - 0.5)) * self.strides
+            y[:, 1::ndim] = (y[:, 1::ndim] * 2.0 + (self.anchors[1] - 0.5)) * self.strides
+            return y
+
+
+class Classify(nn.Module):
+    """YOLOv8 classification head, i.e. x(b,c1,20,20) to x(b,c2)."""
+
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1):  # ch_in, ch_out, kernel, stride, padding, groups
+        super().__init__()
+        c_ = 1280  # efficientnet_b0 size
+        self.conv = Conv(c1, c_, k, s, p, g)
+        self.pool = nn.AdaptiveAvgPool2d(1)  # to x(b,c_,1,1)
+        self.drop = nn.Dropout(p=0.0, inplace=True)
+        self.linear = nn.Linear(c_, c2)  # to x(b,c2)
+
+    def forward(self, x):
+        """Performs a forward pass of the YOLO model on input image data."""
+        if isinstance(x, list):
+            x = torch.cat(x, 1)
+        x = self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
+        return x if self.training else x.softmax(1)
+
+
+class RTDETRDecoder(nn.Module):
+
+    def __init__(
+            self,
+            nc=80,
+            ch=(512, 1024, 2048),
+            hd=256,  # hidden dim
+            nq=300,  # num queries
+            ndp=4,  # num decoder points
+            nh=8,  # num head
+            ndl=6,  # num decoder layers
+            d_ffn=1024,  # dim of feedforward
+            dropout=0.,
+            act=nn.ReLU(),
+            eval_idx=-1,
+            # training args
+            nd=100,  # num denoising
+            label_noise_ratio=0.5,
+            box_noise_scale=1.0,
+            learnt_init_query=False):
+        super().__init__()
+        self.hidden_dim = hd
+        self.nhead = nh
+        self.nl = len(ch)  # num level
+        self.nc = nc
+        self.num_queries = nq
+        self.num_decoder_layers = ndl
+
+        # backbone feature projection
+        self.input_proj = nn.ModuleList(nn.Sequential(nn.Conv2d(x, hd, 1, bias=False), nn.BatchNorm2d(hd)) for x in ch)
+        # NOTE: simplified version but it's not consistent with .pt weights.
+        # self.input_proj = nn.ModuleList(Conv(x, hd, act=False) for x in ch)
+
+        # Transformer module
+        decoder_layer = DeformableTransformerDecoderLayer(hd, nh, d_ffn, dropout, act, self.nl, ndp)
+        self.decoder = DeformableTransformerDecoder(hd, decoder_layer, ndl, eval_idx)
+
+        # denoising part
+        self.denoising_class_embed = nn.Embedding(nc, hd)
+        self.num_denoising = nd
+        self.label_noise_ratio = label_noise_ratio
+        self.box_noise_scale = box_noise_scale
+
+        # decoder embedding
+        self.learnt_init_query = learnt_init_query
+        if learnt_init_query:
+            self.tgt_embed = nn.Embedding(nq, hd)
+        self.query_pos_head = MLP(4, 2 * hd, hd, num_layers=2)
+
+        # encoder head
+        self.enc_output = nn.Sequential(nn.Linear(hd, hd), nn.LayerNorm(hd))
+        self.enc_score_head = nn.Linear(hd, nc)
+        self.enc_bbox_head = MLP(hd, hd, 4, num_layers=3)
+
+        # decoder head
+        self.dec_score_head = nn.ModuleList([nn.Linear(hd, nc) for _ in range(ndl)])
+        self.dec_bbox_head = nn.ModuleList([MLP(hd, hd, 4, num_layers=3) for _ in range(ndl)])
+
+        self._reset_parameters()
+
+    def forward(self, x, batch=None):
+        from ultralytics.vit.utils.ops import get_cdn_group
+
+        # input projection and embedding
+        feats, shapes = self._get_encoder_input(x)
+
+        # prepare denoising training
+        dn_embed, dn_bbox, attn_mask, dn_meta = \
+            get_cdn_group(batch,
+                          self.nc,
+                          self.num_queries,
+                          self.denoising_class_embed.weight,
+                          self.num_denoising,
+                          self.label_noise_ratio,
+                          self.box_noise_scale,
+                          self.training)
+
+        embed, refer_bbox, enc_bboxes, enc_scores = \
+            self._get_decoder_input(feats, shapes, dn_embed, dn_bbox)
+
+        # decoder
+        dec_bboxes, dec_scores = self.decoder(embed,
+                                              refer_bbox,
+                                              feats,
+                                              shapes,
+                                              self.dec_bbox_head,
+                                              self.dec_score_head,
+                                              self.query_pos_head,
+                                              attn_mask=attn_mask)
+        if not self.training:
+            dec_scores = dec_scores.sigmoid_()
+        return dec_bboxes, dec_scores, enc_bboxes, enc_scores, dn_meta
+
+    def _generate_anchors(self, shapes, grid_size=0.05, dtype=torch.float32, device='cpu', eps=1e-2):
+        anchors = []
+        for i, (h, w) in enumerate(shapes):
+            grid_y, grid_x = torch.meshgrid(torch.arange(end=h, dtype=dtype, device=device),
+                                            torch.arange(end=w, dtype=dtype, device=device),
+                                            indexing='ij')
+            grid_xy = torch.stack([grid_x, grid_y], -1)  # (h, w, 2)
+
+            valid_WH = torch.tensor([h, w], dtype=dtype, device=device)
+            grid_xy = (grid_xy.unsqueeze(0) + 0.5) / valid_WH  # (1, h, w, 2)
+            wh = torch.ones_like(grid_xy, dtype=dtype, device=device) * grid_size * (2.0 ** i)
+            anchors.append(torch.cat([grid_xy, wh], -1).view(-1, h * w, 4))  # (1, h*w, 4)
+
+        anchors = torch.cat(anchors, 1)  # (1, h*w*nl, 4)
+        valid_mask = ((anchors > eps) * (anchors < 1 - eps)).all(-1, keepdim=True)  # 1, h*w*nl, 1
+        anchors = torch.log(anchors / (1 - anchors))
+        anchors = torch.where(valid_mask, anchors, torch.inf)
+        return anchors, valid_mask
+
+    def _get_encoder_input(self, x):
+        # get projection features
+        x = [self.input_proj[i](feat) for i, feat in enumerate(x)]
+        # get encoder inputs
+        feats = []
+        shapes = []
+        for feat in x:
+            h, w = feat.shape[2:]
+            # [b, c, h, w] -> [b, h*w, c]
+            feats.append(feat.flatten(2).permute(0, 2, 1))
+            # [nl, 2]
+            shapes.append([h, w])
+
+        # [b, h*w, c]
+        feats = torch.cat(feats, 1)
+        return feats, shapes
+
+    def _get_decoder_input(self, feats, shapes, dn_embed=None, dn_bbox=None):
+        bs = len(feats)
+        # prepare input for decoder
+        anchors, valid_mask = self._generate_anchors(shapes, dtype=feats.dtype, device=feats.device)
+        features = self.enc_output(torch.where(valid_mask, feats, 0))  # bs, h*w, 256
+
+        enc_outputs_scores = self.enc_score_head(features)  # (bs, h*w, nc)
+        # dynamic anchors + static content
+        enc_outputs_bboxes = self.enc_bbox_head(features) + anchors  # (bs, h*w, 4)
+
+        # query selection
+        # (bs, num_queries)
+        topk_ind = torch.topk(enc_outputs_scores.max(-1).values, self.num_queries, dim=1).indices.view(-1)
+        # (bs, num_queries)
+        batch_ind = torch.arange(end=bs, dtype=topk_ind.dtype).unsqueeze(-1).repeat(1, self.num_queries).view(-1)
+
+        # Unsigmoided
+        refer_bbox = enc_outputs_bboxes[batch_ind, topk_ind].view(bs, self.num_queries, -1)
+        # refer_bbox = torch.gather(enc_outputs_bboxes, 1, topk_ind.reshape(bs, self.num_queries).unsqueeze(-1).repeat(1, 1, 4))
+
+        enc_bboxes = refer_bbox.sigmoid()
+        if dn_bbox is not None:
+            refer_bbox = torch.cat([dn_bbox, refer_bbox], 1)
+        if self.training:
+            refer_bbox = refer_bbox.detach()
+        enc_scores = enc_outputs_scores[batch_ind, topk_ind].view(bs, self.num_queries, -1)
+
+        if self.learnt_init_query:
+            embeddings = self.tgt_embed.weight.unsqueeze(0).repeat(bs, 1, 1)
+        else:
+            embeddings = features[batch_ind, topk_ind].view(bs, self.num_queries, -1)
+            if self.training:
+                embeddings = embeddings.detach()
+        if dn_embed is not None:
+            embeddings = torch.cat([dn_embed, embeddings], 1)
+
+        return embeddings, refer_bbox, enc_bboxes, enc_scores
+
+    # TODO
+    def _reset_parameters(self):
+        # class and bbox head init
+        bias_cls = bias_init_with_prob(0.01) / 80 * self.nc
+        # NOTE: the weight initialization in `linear_init_` would cause NaN when training with custom datasets.
+        # linear_init_(self.enc_score_head)
+        constant_(self.enc_score_head.bias, bias_cls)
+        constant_(self.enc_bbox_head.layers[-1].weight, 0.)
+        constant_(self.enc_bbox_head.layers[-1].bias, 0.)
+        for cls_, reg_ in zip(self.dec_score_head, self.dec_bbox_head):
+            # linear_init_(cls_)
+            constant_(cls_.bias, bias_cls)
+            constant_(reg_.layers[-1].weight, 0.)
+            constant_(reg_.layers[-1].bias, 0.)
+
+        linear_init_(self.enc_output[0])
+        xavier_uniform_(self.enc_output[0].weight)
+        if self.learnt_init_query:
+            xavier_uniform_(self.tgt_embed.weight)
+        xavier_uniform_(self.query_pos_head.layers[0].weight)
+        xavier_uniform_(self.query_pos_head.layers[1].weight)
+        for layer in self.input_proj:
+            xavier_uniform_(layer[0].weight)

ultralytics/nn/modules/transformer.py

@@ -0,0 +1,378 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Transformer modules
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import constant_, xavier_uniform_
+
+from .conv import Conv
+from .utils import _get_clones, inverse_sigmoid, multi_scale_deformable_attn_pytorch
+
+__all__ = ('TransformerEncoderLayer', 'TransformerLayer', 'TransformerBlock', 'MLPBlock', 'LayerNorm2d', 'AIFI',
+           'DeformableTransformerDecoder', 'DeformableTransformerDecoderLayer', 'MSDeformAttn', 'MLP')
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Transformer Encoder."""
+
+    def __init__(self, c1, cm=2048, num_heads=8, dropout=0.0, act=nn.GELU(), normalize_before=False):
+        super().__init__()
+        self.ma = nn.MultiheadAttention(c1, num_heads, dropout=dropout, batch_first=True)
+        # Implementation of Feedforward model
+        self.fc1 = nn.Linear(c1, cm)
+        self.fc2 = nn.Linear(cm, c1)
+
+        self.norm1 = nn.LayerNorm(c1)
+        self.norm2 = nn.LayerNorm(c1)
+        self.dropout = nn.Dropout(dropout)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.act = act
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos=None):
+        """Add position embeddings if given."""
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
+        q = k = self.with_pos_embed(src, pos)
+        src2 = self.ma(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.fc2(self.dropout(self.act(self.fc1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+    def forward_pre(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.ma(q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.fc2(self.dropout(self.act(self.fc1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
+        """Forward propagates the input through the encoder module."""
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class AIFI(TransformerEncoderLayer):
+
+    def __init__(self, c1, cm=2048, num_heads=8, dropout=0, act=nn.GELU(), normalize_before=False):
+        super().__init__(c1, cm, num_heads, dropout, act, normalize_before)
+
+    def forward(self, x):
+        c, h, w = x.shape[1:]
+        pos_embed = self.build_2d_sincos_position_embedding(w, h, c)
+        # flatten [B, C, H, W] to [B, HxW, C]
+        x = super().forward(x.flatten(2).permute(0, 2, 1), pos=pos_embed.to(device=x.device, dtype=x.dtype))
+        return x.permute(0, 2, 1).view([-1, c, h, w]).contiguous()
+
+    @staticmethod
+    def build_2d_sincos_position_embedding(w, h, embed_dim=256, temperature=10000.):
+        grid_w = torch.arange(int(w), dtype=torch.float32)
+        grid_h = torch.arange(int(h), dtype=torch.float32)
+        grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing='ij')
+        assert embed_dim % 4 == 0, \
+            'Embed dimension must be divisible by 4 for 2D sin-cos position embedding'
+        pos_dim = embed_dim // 4
+        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
+        omega = 1. / (temperature ** omega)
+
+        out_w = grid_w.flatten()[..., None] @ omega[None]
+        out_h = grid_h.flatten()[..., None] @ omega[None]
+
+        return torch.concat([torch.sin(out_w), torch.cos(out_w),
+                             torch.sin(out_h), torch.cos(out_h)], axis=1)[None, :, :]
+
+
+class TransformerLayer(nn.Module):
+    """Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)."""
+
+    def __init__(self, c, num_heads):
+        """Initializes a self-attention mechanism using linear transformations and multi-head attention."""
+        super().__init__()
+        self.q = nn.Linear(c, c, bias=False)
+        self.k = nn.Linear(c, c, bias=False)
+        self.v = nn.Linear(c, c, bias=False)
+        self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
+        self.fc1 = nn.Linear(c, c, bias=False)
+        self.fc2 = nn.Linear(c, c, bias=False)
+
+    def forward(self, x):
+        """Apply a transformer block to the input x and return the output."""
+        x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
+        x = self.fc2(self.fc1(x)) + x
+        return x
+
+
+class TransformerBlock(nn.Module):
+    """Vision Transformer https://arxiv.org/abs/2010.11929."""
+
+    def __init__(self, c1, c2, num_heads, num_layers):
+        """Initialize a Transformer module with position embedding and specified number of heads and layers."""
+        super().__init__()
+        self.conv = None
+        if c1 != c2:
+            self.conv = Conv(c1, c2)
+        self.linear = nn.Linear(c2, c2)  # learnable position embedding
+        self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
+        self.c2 = c2
+
+    def forward(self, x):
+        """Forward propagates the input through the bottleneck module."""
+        if self.conv is not None:
+            x = self.conv(x)
+        b, _, w, h = x.shape
+        p = x.flatten(2).permute(2, 0, 1)
+        return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
+
+
+class MLPBlock(nn.Module):
+
+    def __init__(self, embedding_dim, mlp_dim, act=nn.GELU):
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin2(self.act(self.lin1(x)))
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
+class LayerNorm2d(nn.Module):
+
+    def __init__(self, num_channels, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x):
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+class MSDeformAttn(nn.Module):
+    """
+    Original Multi-Scale Deformable Attention Module.
+    https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/modules/ms_deform_attn.py
+    """
+
+    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError(f'd_model must be divisible by n_heads, but got {d_model} and {n_heads}')
+        _d_per_head = d_model // n_heads
+        # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
+        assert _d_per_head * n_heads == d_model, '`d_model` must be divisible by `n_heads`'
+
+        self.im2col_step = 64
+
+        self.d_model = d_model
+        self.n_levels = n_levels
+        self.n_heads = n_heads
+        self.n_points = n_points
+
+        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.output_proj = nn.Linear(d_model, d_model)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.sampling_offsets.weight.data, 0.)
+        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(
+            1, self.n_levels, self.n_points, 1)
+        for i in range(self.n_points):
+            grid_init[:, :, i, :] *= i + 1
+        with torch.no_grad():
+            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+        constant_(self.attention_weights.weight.data, 0.)
+        constant_(self.attention_weights.bias.data, 0.)
+        xavier_uniform_(self.value_proj.weight.data)
+        constant_(self.value_proj.bias.data, 0.)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.)
+
+    def forward(self, query, refer_bbox, value, value_shapes, value_mask=None):
+        """
+        https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
+        Args:
+            query (torch.Tensor): [bs, query_length, C]
+            refer_bbox (torch.Tensor): [bs, query_length, n_levels, 2], range in [0, 1], top-left (0,0),
+                bottom-right (1, 1), including padding area
+            value (torch.Tensor): [bs, value_length, C]
+            value_shapes (List): [n_levels, 2], [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
+            value_mask (Tensor): [bs, value_length], True for non-padding elements, False for padding elements
+
+        Returns:
+            output (Tensor): [bs, Length_{query}, C]
+        """
+        bs, len_q = query.shape[:2]
+        len_v = value.shape[1]
+        assert sum(s[0] * s[1] for s in value_shapes) == len_v
+
+        value = self.value_proj(value)
+        if value_mask is not None:
+            value = value.masked_fill(value_mask[..., None], float(0))
+        value = value.view(bs, len_v, self.n_heads, self.d_model // self.n_heads)
+        sampling_offsets = self.sampling_offsets(query).view(bs, len_q, self.n_heads, self.n_levels, self.n_points, 2)
+        attention_weights = self.attention_weights(query).view(bs, len_q, self.n_heads, self.n_levels * self.n_points)
+        attention_weights = F.softmax(attention_weights, -1).view(bs, len_q, self.n_heads, self.n_levels, self.n_points)
+        # N, Len_q, n_heads, n_levels, n_points, 2
+        num_points = refer_bbox.shape[-1]
+        if num_points == 2:
+            offset_normalizer = torch.as_tensor(value_shapes, dtype=query.dtype, device=query.device).flip(-1)
+            add = sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+            sampling_locations = refer_bbox[:, :, None, :, None, :] + add
+        elif num_points == 4:
+            add = sampling_offsets / self.n_points * refer_bbox[:, :, None, :, None, 2:] * 0.5
+            sampling_locations = refer_bbox[:, :, None, :, None, :2] + add
+        else:
+            raise ValueError(f'Last dim of reference_points must be 2 or 4, but got {num_points}.')
+        output = multi_scale_deformable_attn_pytorch(value, value_shapes, sampling_locations, attention_weights)
+        output = self.output_proj(output)
+        return output
+
+
+class DeformableTransformerDecoderLayer(nn.Module):
+    """
+    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
+    https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/deformable_transformer.py
+    """
+
+    def __init__(self, d_model=256, n_heads=8, d_ffn=1024, dropout=0., act=nn.ReLU(), n_levels=4, n_points=4):
+        super().__init__()
+
+        # self attention
+        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
+        self.dropout1 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # cross attention
+        self.cross_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
+        self.dropout2 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # ffn
+        self.linear1 = nn.Linear(d_model, d_ffn)
+        self.act = act
+        self.dropout3 = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ffn, d_model)
+        self.dropout4 = nn.Dropout(dropout)
+        self.norm3 = nn.LayerNorm(d_model)
+
+    @staticmethod
+    def with_pos_embed(tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward_ffn(self, tgt):
+        tgt2 = self.linear2(self.dropout3(self.act(self.linear1(tgt))))
+        tgt = tgt + self.dropout4(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward(self, embed, refer_bbox, feats, shapes, padding_mask=None, attn_mask=None, query_pos=None):
+        # self attention
+        q = k = self.with_pos_embed(embed, query_pos)
+        tgt = self.self_attn(q.transpose(0, 1), k.transpose(0, 1), embed.transpose(0, 1),
+                             attn_mask=attn_mask)[0].transpose(0, 1)
+        embed = embed + self.dropout1(tgt)
+        embed = self.norm1(embed)
+
+        # cross attention
+        tgt = self.cross_attn(self.with_pos_embed(embed, query_pos), refer_bbox.unsqueeze(2), feats, shapes,
+                              padding_mask)
+        embed = embed + self.dropout2(tgt)
+        embed = self.norm2(embed)
+
+        # ffn
+        embed = self.forward_ffn(embed)
+
+        return embed
+
+
+class DeformableTransformerDecoder(nn.Module):
+    """
+    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
+    """
+
+    def __init__(self, hidden_dim, decoder_layer, num_layers, eval_idx=-1):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx
+
+    def forward(
+            self,
+            embed,  # decoder embeddings
+            refer_bbox,  # anchor
+            feats,  # image features
+            shapes,  # feature shapes
+            bbox_head,
+            score_head,
+            pos_mlp,
+            attn_mask=None,
+            padding_mask=None):
+        output = embed
+        dec_bboxes = []
+        dec_cls = []
+        last_refined_bbox = None
+        refer_bbox = refer_bbox.sigmoid()
+        for i, layer in enumerate(self.layers):
+            output = layer(output, refer_bbox, feats, shapes, padding_mask, attn_mask, pos_mlp(refer_bbox))
+
+            # refine bboxes, (bs, num_queries+num_denoising, 4)
+            refined_bbox = torch.sigmoid(bbox_head[i](output) + inverse_sigmoid(refer_bbox))
+
+            if self.training:
+                dec_cls.append(score_head[i](output))
+                if i == 0:
+                    dec_bboxes.append(refined_bbox)
+                else:
+                    dec_bboxes.append(torch.sigmoid(bbox_head[i](output) + inverse_sigmoid(last_refined_bbox)))
+            elif i == self.eval_idx:
+                dec_cls.append(score_head[i](output))
+                dec_bboxes.append(refined_bbox)
+                break
+
+            last_refined_bbox = refined_bbox
+            refer_bbox = refined_bbox.detach() if self.training else refined_bbox
+
+        return torch.stack(dec_bboxes), torch.stack(dec_cls)

ultralytics/nn/modules/utils.py

@@ -0,0 +1,78 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+"""
+Module utils
+"""
+
+import copy
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import uniform_
+
+__all__ = 'multi_scale_deformable_attn_pytorch', 'inverse_sigmoid'
+
+
+def _get_clones(module, n):
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
+
+
+def bias_init_with_prob(prior_prob=0.01):
+    """initialize conv/fc bias value according to a given probability value."""
+    return float(-np.log((1 - prior_prob) / prior_prob))  # return bias_init
+
+
+def linear_init_(module):
+    bound = 1 / math.sqrt(module.weight.shape[0])
+    uniform_(module.weight, -bound, bound)
+    if hasattr(module, 'bias') and module.bias is not None:
+        uniform_(module.bias, -bound, bound)
+
+
+def inverse_sigmoid(x, eps=1e-5):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1 / x2)
+
+
+def multi_scale_deformable_attn_pytorch(value: torch.Tensor, value_spatial_shapes: torch.Tensor,
+                                        sampling_locations: torch.Tensor,
+                                        attention_weights: torch.Tensor) -> torch.Tensor:
+    """
+    Multi-scale deformable attention.
+    https://github.com/IDEA-Research/detrex/blob/main/detrex/layers/multi_scale_deform_attn.py
+    """
+
+    bs, _, num_heads, embed_dims = value.shape
+    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
+    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
+    sampling_grids = 2 * sampling_locations - 1
+    sampling_value_list = []
+    for level, (H_, W_) in enumerate(value_spatial_shapes):
+        # bs, H_*W_, num_heads, embed_dims ->
+        # bs, H_*W_, num_heads*embed_dims ->
+        # bs, num_heads*embed_dims, H_*W_ ->
+        # bs*num_heads, embed_dims, H_, W_
+        value_l_ = (value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_))
+        # bs, num_queries, num_heads, num_points, 2 ->
+        # bs, num_heads, num_queries, num_points, 2 ->
+        # bs*num_heads, num_queries, num_points, 2
+        sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1)
+        # bs*num_heads, embed_dims, num_queries, num_points
+        sampling_value_l_ = F.grid_sample(value_l_,
+                                          sampling_grid_l_,
+                                          mode='bilinear',
+                                          padding_mode='zeros',
+                                          align_corners=False)
+        sampling_value_list.append(sampling_value_l_)
+    # (bs, num_queries, num_heads, num_levels, num_points) ->
+    # (bs, num_heads, num_queries, num_levels, num_points) ->
+    # (bs, num_heads, 1, num_queries, num_levels*num_points)
+    attention_weights = attention_weights.transpose(1, 2).reshape(bs * num_heads, 1, num_queries,
+                                                                  num_levels * num_points)
+    output = ((torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(
+        bs, num_heads * embed_dims, num_queries))
+    return output.transpose(1, 2).contiguous()

ultralytics/nn/tasks.py

@@ -0,0 +1,780 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import contextlib
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from ultralytics.nn.modules import (AIFI, C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x,
+                                    Classify, Concat, Conv, Conv2, ConvTranspose, Detect, DWConv, DWConvTranspose2d,
+                                    Focus, GhostBottleneck, GhostConv, HGBlock, HGStem, Pose, RepC3, RepConv,
+                                    RTDETRDecoder, Segment)
+from ultralytics.yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, colorstr, emojis, yaml_load
+from ultralytics.yolo.utils.checks import check_requirements, check_suffix, check_yaml
+from ultralytics.yolo.utils.loss import v8ClassificationLoss, v8DetectionLoss, v8PoseLoss, v8SegmentationLoss
+from ultralytics.yolo.utils.plotting import feature_visualization
+from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, fuse_deconv_and_bn, initialize_weights,
+                                                intersect_dicts, make_divisible, model_info, scale_img, time_sync)
+
+try:
+    import thop
+except ImportError:
+    thop = None
+
+
+class BaseModel(nn.Module):
+    """
+    The BaseModel class serves as a base class for all the models in the Ultralytics YOLO family.
+    """
+
+    def forward(self, x, *args, **kwargs):
+        """
+        Forward pass of the model on a single scale.
+        Wrapper for `_forward_once` method.
+
+        Args:
+            x (torch.Tensor | dict): The input image tensor or a dict including image tensor and gt labels.
+
+        Returns:
+            (torch.Tensor): The output of the network.
+        """
+        if isinstance(x, dict):  # for cases of training and validating while training.
+            return self.loss(x, *args, **kwargs)
+        return self.predict(x, *args, **kwargs)
+
+    def predict(self, x, profile=False, visualize=False, augment=False):
+        """
+        Perform a forward pass through the network.
+
+        Args:
+            x (torch.Tensor): The input tensor to the model.
+            profile (bool):  Print the computation time of each layer if True, defaults to False.
+            visualize (bool): Save the feature maps of the model if True, defaults to False.
+            augment (bool): Augment image during prediction, defaults to False.
+
+        Returns:
+            (torch.Tensor): The last output of the model.
+        """
+        if augment:
+            return self._predict_augment(x)
+        return self._predict_once(x, profile, visualize)
+
+    def _predict_once(self, x, profile=False, visualize=False):
+        """
+        Perform a forward pass through the network.
+
+        Args:
+            x (torch.Tensor): The input tensor to the model.
+            profile (bool):  Print the computation time of each layer if True, defaults to False.
+            visualize (bool): Save the feature maps of the model if True, defaults to False.
+
+        Returns:
+            (torch.Tensor): The last output of the model.
+        """
+        y, dt = [], []  # outputs
+        for m in self.model:
+            if m.f != -1:  # if not from previous layer
+                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
+            if profile:
+                self._profile_one_layer(m, x, dt)
+            x = m(x)  # run
+            y.append(x if m.i in self.save else None)  # save output
+            if visualize:
+                feature_visualization(x, m.type, m.i, save_dir=visualize)
+        return x
+
+    def _predict_augment(self, x):
+        """Perform augmentations on input image x and return augmented inference."""
+        LOGGER.warning(
+            f'WARNING ⚠️ {self.__class__.__name__} has not supported augment inference yet! Now using single-scale inference instead.'
+        )
+        return self._predict_once(x)
+
+    def _profile_one_layer(self, m, x, dt):
+        """
+        Profile the computation time and FLOPs of a single layer of the model on a given input.
+        Appends the results to the provided list.
+
+        Args:
+            m (nn.Module): The layer to be profiled.
+            x (torch.Tensor): The input data to the layer.
+            dt (list): A list to store the computation time of the layer.
+
+        Returns:
+            None
+        """
+        c = m == self.model[-1]  # is final layer, copy input as inplace fix
+        o = thop.profile(m, inputs=[x.clone() if c else x], verbose=False)[0] / 1E9 * 2 if thop else 0  # FLOPs
+        t = time_sync()
+        for _ in range(10):
+            m(x.clone() if c else x)
+        dt.append((time_sync() - t) * 100)
+        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}')
+        if c:
+            LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")
+
+    def fuse(self, verbose=True):
+        """
+        Fuse the `Conv2d()` and `BatchNorm2d()` layers of the model into a single layer, in order to improve the
+        computation efficiency.
+
+        Returns:
+            (nn.Module): The fused model is returned.
+        """
+        if not self.is_fused():
+            for m in self.model.modules():
+                if isinstance(m, (Conv, Conv2, DWConv)) and hasattr(m, 'bn'):
+                    if isinstance(m, Conv2):
+                        m.fuse_convs()
+                    m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
+                    delattr(m, 'bn')  # remove batchnorm
+                    m.forward = m.forward_fuse  # update forward
+                if isinstance(m, ConvTranspose) and hasattr(m, 'bn'):
+                    m.conv_transpose = fuse_deconv_and_bn(m.conv_transpose, m.bn)
+                    delattr(m, 'bn')  # remove batchnorm
+                    m.forward = m.forward_fuse  # update forward
+                if isinstance(m, RepConv):
+                    m.fuse_convs()
+                    m.forward = m.forward_fuse  # update forward
+            self.info(verbose=verbose)
+
+        return self
+
+    def is_fused(self, thresh=10):
+        """
+        Check if the model has less than a certain threshold of BatchNorm layers.
+
+        Args:
+            thresh (int, optional): The threshold number of BatchNorm layers. Default is 10.
+
+        Returns:
+            (bool): True if the number of BatchNorm layers in the model is less than the threshold, False otherwise.
+        """
+        bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k)  # normalization layers, i.e. BatchNorm2d()
+        return sum(isinstance(v, bn) for v in self.modules()) < thresh  # True if < 'thresh' BatchNorm layers in model
+
+    def info(self, detailed=False, verbose=True, imgsz=640):
+        """
+        Prints model information
+
+        Args:
+            verbose (bool): if True, prints out the model information. Defaults to False
+            imgsz (int): the size of the image that the model will be trained on. Defaults to 640
+        """
+        return model_info(self, detailed=detailed, verbose=verbose, imgsz=imgsz)
+
+    def _apply(self, fn):
+        """
+        `_apply()` is a function that applies a function to all the tensors in the model that are not
+        parameters or registered buffers
+
+        Args:
+            fn: the function to apply to the model
+
+        Returns:
+            A model that is a Detect() object.
+        """
+        self = super()._apply(fn)
+        m = self.model[-1]  # Detect()
+        if isinstance(m, (Detect, Segment)):
+            m.stride = fn(m.stride)
+            m.anchors = fn(m.anchors)
+            m.strides = fn(m.strides)
+        return self
+
+    def load(self, weights, verbose=True):
+        """Load the weights into the model.
+
+        Args:
+            weights (dict | torch.nn.Module): The pre-trained weights to be loaded.
+            verbose (bool, optional): Whether to log the transfer progress. Defaults to True.
+        """
+        model = weights['model'] if isinstance(weights, dict) else weights  # torchvision models are not dicts
+        csd = model.float().state_dict()  # checkpoint state_dict as FP32
+        csd = intersect_dicts(csd, self.state_dict())  # intersect
+        self.load_state_dict(csd, strict=False)  # load
+        if verbose:
+            LOGGER.info(f'Transferred {len(csd)}/{len(self.model.state_dict())} items from pretrained weights')
+
+    def loss(self, batch, preds=None):
+        """
+        Compute loss
+
+        Args:
+            batch (dict): Batch to compute loss on
+            preds (torch.Tensor | List[torch.Tensor]): Predictions.
+        """
+        if not hasattr(self, 'criterion'):
+            self.criterion = self.init_criterion()
+
+        preds = self.forward(batch['img']) if preds is None else preds
+        return self.criterion(preds, batch)
+
+    def init_criterion(self):
+        raise NotImplementedError('compute_loss() needs to be implemented by task heads')
+
+
+class DetectionModel(BaseModel):
+    """YOLOv8 detection model."""
+
+    def __init__(self, cfg='yolov8n.yaml', ch=3, nc=None, verbose=True):  # model, input channels, number of classes
+        super().__init__()
+        self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg)  # cfg dict
+
+        # Define model
+        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels
+        if nc and nc != self.yaml['nc']:
+            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
+            self.yaml['nc'] = nc  # override yaml value
+        self.model, self.save = parse_model(deepcopy(self.yaml), ch=ch, verbose=verbose)  # model, savelist
+        self.names = {i: f'{i}' for i in range(self.yaml['nc'])}  # default names dict
+        self.inplace = self.yaml.get('inplace', True)
+
+        # Build strides
+        m = self.model[-1]  # Detect()
+        if isinstance(m, (Detect, Segment, Pose)):
+            s = 256  # 2x min stride
+            m.inplace = self.inplace
+            forward = lambda x: self.forward(x)[0] if isinstance(m, (Segment, Pose)) else self.forward(x)
+            m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s))])  # forward
+            self.stride = m.stride
+            m.bias_init()  # only run once
+        else:
+            self.stride = torch.Tensor([32])  # default stride for i.e. RTDETR
+
+        # Init weights, biases
+        initialize_weights(self)
+        if verbose:
+            self.info()
+            LOGGER.info('')
+
+    def _predict_augment(self, x):
+        """Perform augmentations on input image x and return augmented inference and train outputs."""
+        img_size = x.shape[-2:]  # height, width
+        s = [1, 0.83, 0.67]  # scales
+        f = [None, 3, None]  # flips (2-ud, 3-lr)
+        y = []  # outputs
+        for si, fi in zip(s, f):
+            xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
+            yi = super().predict(xi)[0]  # forward
+            # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
+            yi = self._descale_pred(yi, fi, si, img_size)
+            y.append(yi)
+        y = self._clip_augmented(y)  # clip augmented tails
+        return torch.cat(y, -1), None  # augmented inference, train
+
+    @staticmethod
+    def _descale_pred(p, flips, scale, img_size, dim=1):
+        """De-scale predictions following augmented inference (inverse operation)."""
+        p[:, :4] /= scale  # de-scale
+        x, y, wh, cls = p.split((1, 1, 2, p.shape[dim] - 4), dim)
+        if flips == 2:
+            y = img_size[0] - y  # de-flip ud
+        elif flips == 3:
+            x = img_size[1] - x  # de-flip lr
+        return torch.cat((x, y, wh, cls), dim)
+
+    def _clip_augmented(self, y):
+        """Clip YOLOv5 augmented inference tails."""
+        nl = self.model[-1].nl  # number of detection layers (P3-P5)
+        g = sum(4 ** x for x in range(nl))  # grid points
+        e = 1  # exclude layer count
+        i = (y[0].shape[-1] // g) * sum(4 ** x for x in range(e))  # indices
+        y[0] = y[0][..., :-i]  # large
+        i = (y[-1].shape[-1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))  # indices
+        y[-1] = y[-1][..., i:]  # small
+        return y
+
+    def init_criterion(self):
+        return v8DetectionLoss(self)
+
+
+class SegmentationModel(DetectionModel):
+    """YOLOv8 segmentation model."""
+
+    def __init__(self, cfg='yolov8n-seg.yaml', ch=3, nc=None, verbose=True):
+        """Initialize YOLOv8 segmentation model with given config and parameters."""
+        super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
+
+    def init_criterion(self):
+        return v8SegmentationLoss(self)
+
+    def _predict_augment(self, x):
+        """Perform augmentations on input image x and return augmented inference."""
+        LOGGER.warning(
+            f'WARNING ⚠️ {self.__class__.__name__} has not supported augment inference yet! Now using single-scale inference instead.'
+        )
+        return self._predict_once(x)
+
+
+class PoseModel(DetectionModel):
+    """YOLOv8 pose model."""
+
+    def __init__(self, cfg='yolov8n-pose.yaml', ch=3, nc=None, data_kpt_shape=(None, None), verbose=True):
+        """Initialize YOLOv8 Pose model."""
+        if not isinstance(cfg, dict):
+            cfg = yaml_model_load(cfg)  # load model YAML
+        if any(data_kpt_shape) and list(data_kpt_shape) != list(cfg['kpt_shape']):
+            LOGGER.info(f"Overriding model.yaml kpt_shape={cfg['kpt_shape']} with kpt_shape={data_kpt_shape}")
+            cfg['kpt_shape'] = data_kpt_shape
+        super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
+
+    def init_criterion(self):
+        return v8PoseLoss(self)
+
+    def _predict_augment(self, x):
+        """Perform augmentations on input image x and return augmented inference."""
+        LOGGER.warning(
+            f'WARNING ⚠️ {self.__class__.__name__} has not supported augment inference yet! Now using single-scale inference instead.'
+        )
+        return self._predict_once(x)
+
+
+class ClassificationModel(BaseModel):
+    """YOLOv8 classification model."""
+
+    def __init__(self,
+                 cfg=None,
+                 model=None,
+                 ch=3,
+                 nc=None,
+                 cutoff=10,
+                 verbose=True):  # yaml, model, channels, number of classes, cutoff index, verbose flag
+        super().__init__()
+        self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg, ch, nc, verbose)
+
+    def _from_detection_model(self, model, nc=1000, cutoff=10):
+        """Create a YOLOv5 classification model from a YOLOv5 detection model."""
+        from ultralytics.nn.autobackend import AutoBackend
+        if isinstance(model, AutoBackend):
+            model = model.model  # unwrap DetectMultiBackend
+        model.model = model.model[:cutoff]  # backbone
+        m = model.model[-1]  # last layer
+        ch = m.conv.in_channels if hasattr(m, 'conv') else m.cv1.conv.in_channels  # ch into module
+        c = Classify(ch, nc)  # Classify()
+        c.i, c.f, c.type = m.i, m.f, 'models.common.Classify'  # index, from, type
+        model.model[-1] = c  # replace
+        self.model = model.model
+        self.stride = model.stride
+        self.save = []
+        self.nc = nc
+
+    def _from_yaml(self, cfg, ch, nc, verbose):
+        """Set YOLOv8 model configurations and define the model architecture."""
+        self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg)  # cfg dict
+
+        # Define model
+        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels
+        if nc and nc != self.yaml['nc']:
+            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
+            self.yaml['nc'] = nc  # override yaml value
+        elif not nc and not self.yaml.get('nc', None):
+            raise ValueError('nc not specified. Must specify nc in model.yaml or function arguments.')
+        self.model, self.save = parse_model(deepcopy(self.yaml), ch=ch, verbose=verbose)  # model, savelist
+        self.stride = torch.Tensor([1])  # no stride constraints
+        self.names = {i: f'{i}' for i in range(self.yaml['nc'])}  # default names dict
+        self.info()
+
+    @staticmethod
+    def reshape_outputs(model, nc):
+        """Update a TorchVision classification model to class count 'n' if required."""
+        name, m = list((model.model if hasattr(model, 'model') else model).named_children())[-1]  # last module
+        if isinstance(m, Classify):  # YOLO Classify() head
+            if m.linear.out_features != nc:
+                m.linear = nn.Linear(m.linear.in_features, nc)
+        elif isinstance(m, nn.Linear):  # ResNet, EfficientNet
+            if m.out_features != nc:
+                setattr(model, name, nn.Linear(m.in_features, nc))
+        elif isinstance(m, nn.Sequential):
+            types = [type(x) for x in m]
+            if nn.Linear in types:
+                i = types.index(nn.Linear)  # nn.Linear index
+                if m[i].out_features != nc:
+                    m[i] = nn.Linear(m[i].in_features, nc)
+            elif nn.Conv2d in types:
+                i = types.index(nn.Conv2d)  # nn.Conv2d index
+                if m[i].out_channels != nc:
+                    m[i] = nn.Conv2d(m[i].in_channels, nc, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None)
+
+    def init_criterion(self):
+        """Compute the classification loss between predictions and true labels."""
+        return v8ClassificationLoss()
+
+
+class RTDETRDetectionModel(DetectionModel):
+
+    def __init__(self, cfg='rtdetr-l.yaml', ch=3, nc=None, verbose=True):
+        super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
+
+    def init_criterion(self):
+        """Compute the classification loss between predictions and true labels."""
+        from ultralytics.vit.utils.loss import RTDETRDetectionLoss
+
+        return RTDETRDetectionLoss(nc=self.nc, use_vfl=True)
+
+    def loss(self, batch, preds=None):
+        if not hasattr(self, 'criterion'):
+            self.criterion = self.init_criterion()
+
+        img = batch['img']
+        # NOTE: preprocess gt_bbox and gt_labels to list.
+        bs = len(img)
+        batch_idx = batch['batch_idx']
+        gt_groups = [(batch_idx == i).sum().item() for i in range(bs)]
+        targets = {
+            'cls': batch['cls'].to(img.device, dtype=torch.long).view(-1),
+            'bboxes': batch['bboxes'].to(device=img.device),
+            'batch_idx': batch_idx.to(img.device, dtype=torch.long).view(-1),
+            'gt_groups': gt_groups}
+
+        preds = self.predict(img, batch=targets) if preds is None else preds
+        dec_bboxes, dec_scores, enc_bboxes, enc_scores, dn_meta = preds
+        if dn_meta is None:
+            dn_bboxes, dn_scores = None, None
+        else:
+            dn_bboxes, dec_bboxes = torch.split(dec_bboxes, dn_meta['dn_num_split'], dim=2)
+            dn_scores, dec_scores = torch.split(dec_scores, dn_meta['dn_num_split'], dim=2)
+
+        dec_bboxes = torch.cat([enc_bboxes.unsqueeze(0), dec_bboxes])  # (7, bs, 300, 4)
+        dec_scores = torch.cat([enc_scores.unsqueeze(0), dec_scores])
+
+        loss = self.criterion((dec_bboxes, dec_scores),
+                              targets,
+                              dn_bboxes=dn_bboxes,
+                              dn_scores=dn_scores,
+                              dn_meta=dn_meta)
+        # NOTE: There are like 12 losses in RTDETR, backward with all losses but only show the main three losses.
+        return sum(loss.values()), torch.as_tensor([loss[k].detach() for k in ['loss_giou', 'loss_class', 'loss_bbox']],
+                                                   device=img.device)
+
+    def predict(self, x, profile=False, visualize=False, batch=None, augment=False):
+        """
+        Perform a forward pass through the network.
+
+        Args:
+            x (torch.Tensor): The input tensor to the model
+            profile (bool):  Print the computation time of each layer if True, defaults to False.
+            visualize (bool): Save the feature maps of the model if True, defaults to False
+            batch (dict): A dict including gt boxes and labels from dataloader.
+
+        Returns:
+            (torch.Tensor): The last output of the model.
+        """
+        y, dt = [], []  # outputs
+        for m in self.model[:-1]:  # except the head part
+            if m.f != -1:  # if not from previous layer
+                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
+            if profile:
+                self._profile_one_layer(m, x, dt)
+            x = m(x)  # run
+            y.append(x if m.i in self.save else None)  # save output
+            if visualize:
+                feature_visualization(x, m.type, m.i, save_dir=visualize)
+        head = self.model[-1]
+        x = head([y[j] for j in head.f], batch)  # head inference
+        return x
+
+
+class Ensemble(nn.ModuleList):
+    """Ensemble of models."""
+
+    def __init__(self):
+        """Initialize an ensemble of models."""
+        super().__init__()
+
+    def forward(self, x, augment=False, profile=False, visualize=False):
+        """Function generates the YOLOv5 network's final layer."""
+        y = [module(x, augment, profile, visualize)[0] for module in self]
+        # y = torch.stack(y).max(0)[0]  # max ensemble
+        # y = torch.stack(y).mean(0)  # mean ensemble
+        y = torch.cat(y, 2)  # nms ensemble, y shape(B, HW, C)
+        return y, None  # inference, train output
+
+
+# Functions ------------------------------------------------------------------------------------------------------------
+
+
+def torch_safe_load(weight):
+    """
+    This function attempts to load a PyTorch model with the torch.load() function. If a ModuleNotFoundError is raised,
+    it catches the error, logs a warning message, and attempts to install the missing module via the
+    check_requirements() function. After installation, the function again attempts to load the model using torch.load().
+
+    Args:
+        weight (str): The file path of the PyTorch model.
+
+    Returns:
+        (dict): The loaded PyTorch model.
+    """
+    from ultralytics.yolo.utils.downloads import attempt_download_asset
+
+    check_suffix(file=weight, suffix='.pt')
+    file = attempt_download_asset(weight)  # search online if missing locally
+    try:
+        return torch.load(file, map_location='cpu'), file  # load
+    except ModuleNotFoundError as e:  # e.name is missing module name
+        if e.name == 'models':
+            raise TypeError(
+                emojis(f'ERROR ❌️ {weight} appears to be an Ultralytics YOLOv5 model originally trained '
+                       f'with https://github.com/ultralytics/yolov5.\nThis model is NOT forwards compatible with '
+                       f'YOLOv8 at https://github.com/ultralytics/ultralytics.'
+                       f"\nRecommend fixes are to train a new model using the latest 'ultralytics' package or to "
+                       f"run a command with an official YOLOv8 model, i.e. 'yolo predict model=yolov8n.pt'")) from e
+        LOGGER.warning(f"WARNING ⚠️ {weight} appears to require '{e.name}', which is not in ultralytics requirements."
+                       f"\nAutoInstall will run now for '{e.name}' but this feature will be removed in the future."
+                       f"\nRecommend fixes are to train a new model using the latest 'ultralytics' package or to "
+                       f"run a command with an official YOLOv8 model, i.e. 'yolo predict model=yolov8n.pt'")
+        check_requirements(e.name)  # install missing module
+
+        return torch.load(file, map_location='cpu'), file  # load
+
+
+def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
+    """Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a."""
+
+    ensemble = Ensemble()
+    for w in weights if isinstance(weights, list) else [weights]:
+        ckpt, w = torch_safe_load(w)  # load ckpt
+        args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} if 'train_args' in ckpt else None  # combined args
+        model = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model
+
+        # Model compatibility updates
+        model.args = args  # attach args to model
+        model.pt_path = w  # attach *.pt file path to model
+        model.task = guess_model_task(model)
+        if not hasattr(model, 'stride'):
+            model.stride = torch.tensor([32.])
+
+        # Append
+        ensemble.append(model.fuse().eval() if fuse and hasattr(model, 'fuse') else model.eval())  # model in eval mode
+
+    # Module compatibility updates
+    for m in ensemble.modules():
+        t = type(m)
+        if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Segment):
+            m.inplace = inplace  # torch 1.7.0 compatibility
+        elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
+            m.recompute_scale_factor = None  # torch 1.11.0 compatibility
+
+    # Return model
+    if len(ensemble) == 1:
+        return ensemble[-1]
+
+    # Return ensemble
+    LOGGER.info(f'Ensemble created with {weights}\n')
+    for k in 'names', 'nc', 'yaml':
+        setattr(ensemble, k, getattr(ensemble[0], k))
+    ensemble.stride = ensemble[torch.argmax(torch.tensor([m.stride.max() for m in ensemble])).int()].stride
+    assert all(ensemble[0].nc == m.nc for m in ensemble), f'Models differ in class counts {[m.nc for m in ensemble]}'
+    return ensemble
+
+
+def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
+    """Loads a single model weights."""
+    ckpt, weight = torch_safe_load(weight)  # load ckpt
+    args = {**DEFAULT_CFG_DICT, **(ckpt.get('train_args', {}))}  # combine model and default args, preferring model args
+    model = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model
+
+    # Model compatibility updates
+    model.args = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS}  # attach args to model
+    model.pt_path = weight  # attach *.pt file path to model
+    model.task = guess_model_task(model)
+    if not hasattr(model, 'stride'):
+        model.stride = torch.tensor([32.])
+
+    model = model.fuse().eval() if fuse and hasattr(model, 'fuse') else model.eval()  # model in eval mode
+
+    # Module compatibility updates
+    for m in model.modules():
+        t = type(m)
+        if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Segment):
+            m.inplace = inplace  # torch 1.7.0 compatibility
+        elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
+            m.recompute_scale_factor = None  # torch 1.11.0 compatibility
+
+    # Return model and ckpt
+    return model, ckpt
+
+
+def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)
+    # Parse a YOLO model.yaml dictionary into a PyTorch model
+    import ast
+
+    # Args
+    max_channels = float('inf')
+    nc, act, scales = (d.get(x) for x in ('nc', 'activation', 'scales'))
+    depth, width, kpt_shape = (d.get(x, 1.0) for x in ('depth_multiple', 'width_multiple', 'kpt_shape'))
+    if scales:
+        scale = d.get('scale')
+        if not scale:
+            scale = tuple(scales.keys())[0]
+            LOGGER.warning(f"WARNING ⚠️ no model scale passed. Assuming scale='{scale}'.")
+        depth, width, max_channels = scales[scale]
+
+    if act:
+        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
+        if verbose:
+            LOGGER.info(f"{colorstr('activation:')} {act}")  # print
+
+    if verbose:
+        LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10}  {'module':<45}{'arguments':<30}")
+    ch = [ch]
+    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
+    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
+        m = getattr(torch.nn, m[3:]) if 'nn.' in m else globals()[m]  # get module
+        for j, a in enumerate(args):
+            if isinstance(a, str):
+                with contextlib.suppress(ValueError):
+                    args[j] = locals()[a] if a in locals() else ast.literal_eval(a)
+
+        n = n_ = max(round(n * depth), 1) if n > 1 else n  # depth gain
+        if m in (Classify, Conv, ConvTranspose, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, Focus,
+                 BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x, RepC3):
+            c1, c2 = ch[f], args[0]
+            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
+                c2 = make_divisible(min(c2, max_channels) * width, 8)
+
+            args = [c1, c2, *args[1:]]
+            if m in (BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, C3x, RepC3):
+                args.insert(2, n)  # number of repeats
+                n = 1
+        elif m is AIFI:
+            args = [ch[f], *args]
+        elif m in (HGStem, HGBlock):
+            c1, cm, c2 = ch[f], args[0], args[1]
+            args = [c1, cm, c2, *args[2:]]
+            if m is HGBlock:
+                args.insert(4, n)  # number of repeats
+                n = 1
+
+        elif m is nn.BatchNorm2d:
+            args = [ch[f]]
+        elif m is Concat:
+            c2 = sum(ch[x] for x in f)
+        elif m in (Detect, Segment, Pose, RTDETRDecoder):
+            args.append([ch[x] for x in f])
+            if m is Segment:
+                args[2] = make_divisible(min(args[2], max_channels) * width, 8)
+        else:
+            c2 = ch[f]
+
+        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
+        t = str(m)[8:-2].replace('__main__.', '')  # module type
+        m.np = sum(x.numel() for x in m_.parameters())  # number params
+        m_.i, m_.f, m_.type = i, f, t  # attach index, 'from' index, type
+        if verbose:
+            LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f}  {t:<45}{str(args):<30}')  # print
+        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
+        layers.append(m_)
+        if i == 0:
+            ch = []
+        ch.append(c2)
+    return nn.Sequential(*layers), sorted(save)
+
+
+def yaml_model_load(path):
+    """Load a YOLOv8 model from a YAML file."""
+    import re
+
+    path = Path(path)
+    if path.stem in (f'yolov{d}{x}6' for x in 'nsmlx' for d in (5, 8)):
+        new_stem = re.sub(r'(\d+)([nslmx])6(.+)?$', r'\1\2-p6\3', path.stem)
+        LOGGER.warning(f'WARNING ⚠️ Ultralytics YOLO P6 models now use -p6 suffix. Renaming {path.stem} to {new_stem}.')
+        path = path.with_stem(new_stem)
+
+    unified_path = re.sub(r'(\d+)([nslmx])(.+)?$', r'\1\3', str(path))  # i.e. yolov8x.yaml -> yolov8.yaml
+    yaml_file = check_yaml(unified_path, hard=False) or check_yaml(path)
+    d = yaml_load(yaml_file)  # model dict
+    d['scale'] = guess_model_scale(path)
+    d['yaml_file'] = str(path)
+    return d
+
+
+def guess_model_scale(model_path):
+    """
+    Takes a path to a YOLO model's YAML file as input and extracts the size character of the model's scale.
+    The function uses regular expression matching to find the pattern of the model scale in the YAML file name,
+    which is denoted by n, s, m, l, or x. The function returns the size character of the model scale as a string.
+
+    Args:
+        model_path (str | Path): The path to the YOLO model's YAML file.
+
+    Returns:
+        (str): The size character of the model's scale, which can be n, s, m, l, or x.
+    """
+    with contextlib.suppress(AttributeError):
+        import re
+        return re.search(r'yolov\d+([nslmx])', Path(model_path).stem).group(1)  # n, s, m, l, or x
+    return ''
+
+
+def guess_model_task(model):
+    """
+    Guess the task of a PyTorch model from its architecture or configuration.
+
+    Args:
+        model (nn.Module | dict): PyTorch model or model configuration in YAML format.
+
+    Returns:
+        (str): Task of the model ('detect', 'segment', 'classify', 'pose').
+
+    Raises:
+        SyntaxError: If the task of the model could not be determined.
+    """
+
+    def cfg2task(cfg):
+        """Guess from YAML dictionary."""
+        m = cfg['head'][-1][-2].lower()  # output module name
+        if m in ('classify', 'classifier', 'cls', 'fc'):
+            return 'classify'
+        if m == 'detect':
+            return 'detect'
+        if m == 'segment':
+            return 'segment'
+        if m == 'pose':
+            return 'pose'
+
+    # Guess from model cfg
+    if isinstance(model, dict):
+        with contextlib.suppress(Exception):
+            return cfg2task(model)
+
+    # Guess from PyTorch model
+    if isinstance(model, nn.Module):  # PyTorch model
+        for x in 'model.args', 'model.model.args', 'model.model.model.args':
+            with contextlib.suppress(Exception):
+                return eval(x)['task']
+        for x in 'model.yaml', 'model.model.yaml', 'model.model.model.yaml':
+            with contextlib.suppress(Exception):
+                return cfg2task(eval(x))
+
+        for m in model.modules():
+            if isinstance(m, Detect):
+                return 'detect'
+            elif isinstance(m, Segment):
+                return 'segment'
+            elif isinstance(m, Classify):
+                return 'classify'
+            elif isinstance(m, Pose):
+                return 'pose'
+
+    # Guess from model filename
+    if isinstance(model, (str, Path)):
+        model = Path(model)
+        if '-seg' in model.stem or 'segment' in model.parts:
+            return 'segment'
+        elif '-cls' in model.stem or 'classify' in model.parts:
+            return 'classify'
+        elif '-pose' in model.stem or 'pose' in model.parts:
+            return 'pose'
+        elif 'detect' in model.parts:
+            return 'detect'
+
+    # Unable to determine task from model
+    LOGGER.warning("WARNING ⚠️ Unable to automatically guess model task, assuming 'task=detect'. "
+                   "Explicitly define task for your model, i.e. 'task=detect', 'segment', 'classify', or 'pose'.")
+    return 'detect'  # assume detect