├── .gitignore ├── README.md ├── image_classification ├── README.md ├── data │ ├── cat.jpg │ └── labels_mobilenet_quant_v1_224.txt ├── test_camera.py └── test_image.py ├── object_detection ├── README.md ├── data │ ├── bus.jpg │ └── coco_labels.txt ├── test_camera.py └── test_image.py ├── pose_detection ├── README.md ├── data │ └── person.jpg ├── test_camera.py └── test_image.py └── test.PNG /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # TF Models 132 | *.tflite 133 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ai-tflite-opencv 2 | This repo contains codes that use Tensorflow (or Tensorflow Lite) + OpenCV to do object detection, classification, and Pose detection. 3 | 4 | 5 | ## Pre-requisite 6 | 7 | - Installed Tensorflow (or Tensorflow Lite) 8 | - Installed OpenCV 9 | - Has camera (Optional) 10 | 11 | ## Download the Models 12 | 13 | The .tflite models are not included, you will need to download the models before you run the code. 14 | 15 | ## Test Platform 16 | 17 | The codes are tested successfully in below platform. 18 | 19 | - ### Windows 10 / Ubuntu 18.04 20 | 21 | Tensorflow 2.1 + OpenCV 4.2 22 | 23 | - ### Raspberry PI 4 (4G Ram) - Raspbian Buster 24 | 25 | Tensorflow Lite 2.1 + OpenCV 4.1 26 | 27 | 28 | ## Example Run 29 | 30 | Image Classification with camera running on Raspberry PI 4: 31 | 32 | -------------------------------------------------------------------------------- /image_classification/README.md: -------------------------------------------------------------------------------- 1 | # Image Classification 2 | 3 | Use Tensorflow (Lite) + OpenCV to do image classification. 4 | 5 | ## Download the Model 6 | 7 | The tested model can be downloaded from: 8 | 9 | https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_quant_and_labels.zip -------------------------------------------------------------------------------- /image_classification/data/cat.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aiden-dai/ai-tflite-opencv/81971855d81c101f50b2abeaa683577ceca75bf5/image_classification/data/cat.jpg -------------------------------------------------------------------------------- /image_classification/data/labels_mobilenet_quant_v1_224.txt: -------------------------------------------------------------------------------- 1 | background 2 | tench 3 | goldfish 4 | great white shark 5 | tiger shark 6 | hammerhead 7 | electric ray 8 | stingray 9 | cock 10 | hen 11 | ostrich 12 | brambling 13 | goldfinch 14 | house finch 15 | junco 16 | indigo bunting 17 | robin 18 | bulbul 19 | jay 20 | magpie 21 | chickadee 22 | water ouzel 23 | kite 24 | bald eagle 25 | vulture 26 | great grey owl 27 | European fire salamander 28 | common newt 29 | eft 30 | spotted salamander 31 | axolotl 32 | bullfrog 33 | tree frog 34 | tailed frog 35 | loggerhead 36 | leatherback turtle 37 | mud turtle 38 | terrapin 39 | box turtle 40 | banded gecko 41 | common iguana 42 | American chameleon 43 | whiptail 44 | agama 45 | frilled lizard 46 | alligator lizard 47 | Gila monster 48 | green lizard 49 | African chameleon 50 | Komodo dragon 51 | African crocodile 52 | American alligator 53 | triceratops 54 | thunder snake 55 | ringneck snake 56 | hognose snake 57 | green snake 58 | king snake 59 | garter snake 60 | water snake 61 | vine snake 62 | night snake 63 | boa constrictor 64 | rock python 65 | Indian cobra 66 | green mamba 67 | sea snake 68 | horned viper 69 | diamondback 70 | sidewinder 71 | trilobite 72 | harvestman 73 | scorpion 74 | black and gold garden spider 75 | barn spider 76 | garden spider 77 | black widow 78 | tarantula 79 | wolf spider 80 | tick 81 | centipede 82 | black grouse 83 | ptarmigan 84 | ruffed grouse 85 | prairie chicken 86 | peacock 87 | quail 88 | partridge 89 | African grey 90 | macaw 91 | sulphur-crested cockatoo 92 | lorikeet 93 | coucal 94 | bee eater 95 | hornbill 96 | hummingbird 97 | jacamar 98 | toucan 99 | drake 100 | red-breasted merganser 101 | goose 102 | black swan 103 | tusker 104 | echidna 105 | platypus 106 | wallaby 107 | koala 108 | wombat 109 | jellyfish 110 | sea anemone 111 | brain coral 112 | flatworm 113 | nematode 114 | conch 115 | snail 116 | slug 117 | sea slug 118 | chiton 119 | chambered nautilus 120 | Dungeness crab 121 | rock crab 122 | fiddler crab 123 | king crab 124 | American lobster 125 | spiny lobster 126 | crayfish 127 | hermit crab 128 | isopod 129 | white stork 130 | black stork 131 | spoonbill 132 | flamingo 133 | little blue heron 134 | American egret 135 | bittern 136 | crane 137 | limpkin 138 | European gallinule 139 | American coot 140 | bustard 141 | ruddy turnstone 142 | red-backed sandpiper 143 | redshank 144 | dowitcher 145 | oystercatcher 146 | pelican 147 | king penguin 148 | albatross 149 | grey whale 150 | killer whale 151 | dugong 152 | sea lion 153 | Chihuahua 154 | Japanese spaniel 155 | Maltese dog 156 | Pekinese 157 | Shih-Tzu 158 | Blenheim spaniel 159 | papillon 160 | toy terrier 161 | Rhodesian ridgeback 162 | Afghan hound 163 | basset 164 | beagle 165 | bloodhound 166 | bluetick 167 | black-and-tan coonhound 168 | Walker hound 169 | English foxhound 170 | redbone 171 | borzoi 172 | Irish wolfhound 173 | Italian greyhound 174 | whippet 175 | Ibizan hound 176 | Norwegian elkhound 177 | otterhound 178 | Saluki 179 | Scottish deerhound 180 | Weimaraner 181 | Staffordshire bullterrier 182 | American Staffordshire terrier 183 | Bedlington terrier 184 | Border terrier 185 | Kerry blue terrier 186 | Irish terrier 187 | Norfolk terrier 188 | Norwich terrier 189 | Yorkshire terrier 190 | wire-haired fox terrier 191 | Lakeland terrier 192 | Sealyham terrier 193 | Airedale 194 | cairn 195 | Australian terrier 196 | Dandie Dinmont 197 | Boston bull 198 | miniature schnauzer 199 | giant schnauzer 200 | standard schnauzer 201 | Scotch terrier 202 | Tibetan terrier 203 | silky terrier 204 | soft-coated wheaten terrier 205 | West Highland white terrier 206 | Lhasa 207 | flat-coated retriever 208 | curly-coated retriever 209 | golden retriever 210 | Labrador retriever 211 | Chesapeake Bay retriever 212 | German short-haired pointer 213 | vizsla 214 | English setter 215 | Irish setter 216 | Gordon setter 217 | Brittany spaniel 218 | clumber 219 | English springer 220 | Welsh springer spaniel 221 | cocker spaniel 222 | Sussex spaniel 223 | Irish water spaniel 224 | kuvasz 225 | schipperke 226 | groenendael 227 | malinois 228 | briard 229 | kelpie 230 | komondor 231 | Old English sheepdog 232 | Shetland sheepdog 233 | collie 234 | Border collie 235 | Bouvier des Flandres 236 | Rottweiler 237 | German shepherd 238 | Doberman 239 | miniature pinscher 240 | Greater Swiss Mountain dog 241 | Bernese mountain dog 242 | Appenzeller 243 | EntleBucher 244 | boxer 245 | bull mastiff 246 | Tibetan mastiff 247 | French bulldog 248 | Great Dane 249 | Saint Bernard 250 | Eskimo dog 251 | malamute 252 | Siberian husky 253 | dalmatian 254 | affenpinscher 255 | basenji 256 | pug 257 | Leonberg 258 | Newfoundland 259 | Great Pyrenees 260 | Samoyed 261 | Pomeranian 262 | chow 263 | keeshond 264 | Brabancon griffon 265 | Pembroke 266 | Cardigan 267 | toy poodle 268 | miniature poodle 269 | standard poodle 270 | Mexican hairless 271 | timber wolf 272 | white wolf 273 | red wolf 274 | coyote 275 | dingo 276 | dhole 277 | African hunting dog 278 | hyena 279 | red fox 280 | kit fox 281 | Arctic fox 282 | grey fox 283 | tabby 284 | tiger cat 285 | Persian cat 286 | Siamese cat 287 | Egyptian cat 288 | cougar 289 | lynx 290 | leopard 291 | snow leopard 292 | jaguar 293 | lion 294 | tiger 295 | cheetah 296 | brown bear 297 | American black bear 298 | ice bear 299 | sloth bear 300 | mongoose 301 | meerkat 302 | tiger beetle 303 | ladybug 304 | ground beetle 305 | long-horned beetle 306 | leaf beetle 307 | dung beetle 308 | rhinoceros beetle 309 | weevil 310 | fly 311 | bee 312 | ant 313 | grasshopper 314 | cricket 315 | walking stick 316 | cockroach 317 | mantis 318 | cicada 319 | leafhopper 320 | lacewing 321 | dragonfly 322 | damselfly 323 | admiral 324 | ringlet 325 | monarch 326 | cabbage butterfly 327 | sulphur butterfly 328 | lycaenid 329 | starfish 330 | sea urchin 331 | sea cucumber 332 | wood rabbit 333 | hare 334 | Angora 335 | hamster 336 | porcupine 337 | fox squirrel 338 | marmot 339 | beaver 340 | guinea pig 341 | sorrel 342 | zebra 343 | hog 344 | wild boar 345 | warthog 346 | hippopotamus 347 | ox 348 | water buffalo 349 | bison 350 | ram 351 | bighorn 352 | ibex 353 | hartebeest 354 | impala 355 | gazelle 356 | Arabian camel 357 | llama 358 | weasel 359 | mink 360 | polecat 361 | black-footed ferret 362 | otter 363 | skunk 364 | badger 365 | armadillo 366 | three-toed sloth 367 | orangutan 368 | gorilla 369 | chimpanzee 370 | gibbon 371 | siamang 372 | guenon 373 | patas 374 | baboon 375 | macaque 376 | langur 377 | colobus 378 | proboscis monkey 379 | marmoset 380 | capuchin 381 | howler monkey 382 | titi 383 | spider monkey 384 | squirrel monkey 385 | Madagascar cat 386 | indri 387 | Indian elephant 388 | African elephant 389 | lesser panda 390 | giant panda 391 | barracouta 392 | eel 393 | coho 394 | rock beauty 395 | anemone fish 396 | sturgeon 397 | gar 398 | lionfish 399 | puffer 400 | abacus 401 | abaya 402 | academic gown 403 | accordion 404 | acoustic guitar 405 | aircraft carrier 406 | airliner 407 | airship 408 | altar 409 | ambulance 410 | amphibian 411 | analog clock 412 | apiary 413 | apron 414 | ashcan 415 | assault rifle 416 | backpack 417 | bakery 418 | balance beam 419 | balloon 420 | ballpoint 421 | Band Aid 422 | banjo 423 | bannister 424 | barbell 425 | barber chair 426 | barbershop 427 | barn 428 | barometer 429 | barrel 430 | barrow 431 | baseball 432 | basketball 433 | bassinet 434 | bassoon 435 | bathing cap 436 | bath towel 437 | bathtub 438 | beach wagon 439 | beacon 440 | beaker 441 | bearskin 442 | beer bottle 443 | beer glass 444 | bell cote 445 | bib 446 | bicycle-built-for-two 447 | bikini 448 | binder 449 | binoculars 450 | birdhouse 451 | boathouse 452 | bobsled 453 | bolo tie 454 | bonnet 455 | bookcase 456 | bookshop 457 | bottlecap 458 | bow 459 | bow tie 460 | brass 461 | brassiere 462 | breakwater 463 | breastplate 464 | broom 465 | bucket 466 | buckle 467 | bulletproof vest 468 | bullet train 469 | butcher shop 470 | cab 471 | caldron 472 | candle 473 | cannon 474 | canoe 475 | can opener 476 | cardigan 477 | car mirror 478 | carousel 479 | carpenter's kit 480 | carton 481 | car wheel 482 | cash machine 483 | cassette 484 | cassette player 485 | castle 486 | catamaran 487 | CD player 488 | cello 489 | cellular telephone 490 | chain 491 | chainlink fence 492 | chain mail 493 | chain saw 494 | chest 495 | chiffonier 496 | chime 497 | china cabinet 498 | Christmas stocking 499 | church 500 | cinema 501 | cleaver 502 | cliff dwelling 503 | cloak 504 | clog 505 | cocktail shaker 506 | coffee mug 507 | coffeepot 508 | coil 509 | combination lock 510 | computer keyboard 511 | confectionery 512 | container ship 513 | convertible 514 | corkscrew 515 | cornet 516 | cowboy boot 517 | cowboy hat 518 | cradle 519 | crane 520 | crash helmet 521 | crate 522 | crib 523 | Crock Pot 524 | croquet ball 525 | crutch 526 | cuirass 527 | dam 528 | desk 529 | desktop computer 530 | dial telephone 531 | diaper 532 | digital clock 533 | digital watch 534 | dining table 535 | dishrag 536 | dishwasher 537 | disk brake 538 | dock 539 | dogsled 540 | dome 541 | doormat 542 | drilling platform 543 | drum 544 | drumstick 545 | dumbbell 546 | Dutch oven 547 | electric fan 548 | electric guitar 549 | electric locomotive 550 | entertainment center 551 | envelope 552 | espresso maker 553 | face powder 554 | feather boa 555 | file 556 | fireboat 557 | fire engine 558 | fire screen 559 | flagpole 560 | flute 561 | folding chair 562 | football helmet 563 | forklift 564 | fountain 565 | fountain pen 566 | four-poster 567 | freight car 568 | French horn 569 | frying pan 570 | fur coat 571 | garbage truck 572 | gasmask 573 | gas pump 574 | goblet 575 | go-kart 576 | golf ball 577 | golfcart 578 | gondola 579 | gong 580 | gown 581 | grand piano 582 | greenhouse 583 | grille 584 | grocery store 585 | guillotine 586 | hair slide 587 | hair spray 588 | half track 589 | hammer 590 | hamper 591 | hand blower 592 | hand-held computer 593 | handkerchief 594 | hard disc 595 | harmonica 596 | harp 597 | harvester 598 | hatchet 599 | holster 600 | home theater 601 | honeycomb 602 | hook 603 | hoopskirt 604 | horizontal bar 605 | horse cart 606 | hourglass 607 | iPod 608 | iron 609 | jack-o'-lantern 610 | jean 611 | jeep 612 | jersey 613 | jigsaw puzzle 614 | jinrikisha 615 | joystick 616 | kimono 617 | knee pad 618 | knot 619 | lab coat 620 | ladle 621 | lampshade 622 | laptop 623 | lawn mower 624 | lens cap 625 | letter opener 626 | library 627 | lifeboat 628 | lighter 629 | limousine 630 | liner 631 | lipstick 632 | Loafer 633 | lotion 634 | loudspeaker 635 | loupe 636 | lumbermill 637 | magnetic compass 638 | mailbag 639 | mailbox 640 | maillot 641 | maillot 642 | manhole cover 643 | maraca 644 | marimba 645 | mask 646 | matchstick 647 | maypole 648 | maze 649 | measuring cup 650 | medicine chest 651 | megalith 652 | microphone 653 | microwave 654 | military uniform 655 | milk can 656 | minibus 657 | miniskirt 658 | minivan 659 | missile 660 | mitten 661 | mixing bowl 662 | mobile home 663 | Model T 664 | modem 665 | monastery 666 | monitor 667 | moped 668 | mortar 669 | mortarboard 670 | mosque 671 | mosquito net 672 | motor scooter 673 | mountain bike 674 | mountain tent 675 | mouse 676 | mousetrap 677 | moving van 678 | muzzle 679 | nail 680 | neck brace 681 | necklace 682 | nipple 683 | notebook 684 | obelisk 685 | oboe 686 | ocarina 687 | odometer 688 | oil filter 689 | organ 690 | oscilloscope 691 | overskirt 692 | oxcart 693 | oxygen mask 694 | packet 695 | paddle 696 | paddlewheel 697 | padlock 698 | paintbrush 699 | pajama 700 | palace 701 | panpipe 702 | paper towel 703 | parachute 704 | parallel bars 705 | park bench 706 | parking meter 707 | passenger car 708 | patio 709 | pay-phone 710 | pedestal 711 | pencil box 712 | pencil sharpener 713 | perfume 714 | Petri dish 715 | photocopier 716 | pick 717 | pickelhaube 718 | picket fence 719 | pickup 720 | pier 721 | piggy bank 722 | pill bottle 723 | pillow 724 | ping-pong ball 725 | pinwheel 726 | pirate 727 | pitcher 728 | plane 729 | planetarium 730 | plastic bag 731 | plate rack 732 | plow 733 | plunger 734 | Polaroid camera 735 | pole 736 | police van 737 | poncho 738 | pool table 739 | pop bottle 740 | pot 741 | potter's wheel 742 | power drill 743 | prayer rug 744 | printer 745 | prison 746 | projectile 747 | projector 748 | puck 749 | punching bag 750 | purse 751 | quill 752 | quilt 753 | racer 754 | racket 755 | radiator 756 | radio 757 | radio telescope 758 | rain barrel 759 | recreational vehicle 760 | reel 761 | reflex camera 762 | refrigerator 763 | remote control 764 | restaurant 765 | revolver 766 | rifle 767 | rocking chair 768 | rotisserie 769 | rubber eraser 770 | rugby ball 771 | rule 772 | running shoe 773 | safe 774 | safety pin 775 | saltshaker 776 | sandal 777 | sarong 778 | sax 779 | scabbard 780 | scale 781 | school bus 782 | schooner 783 | scoreboard 784 | screen 785 | screw 786 | screwdriver 787 | seat belt 788 | sewing machine 789 | shield 790 | shoe shop 791 | shoji 792 | shopping basket 793 | shopping cart 794 | shovel 795 | shower cap 796 | shower curtain 797 | ski 798 | ski mask 799 | sleeping bag 800 | slide rule 801 | sliding door 802 | slot 803 | snorkel 804 | snowmobile 805 | snowplow 806 | soap dispenser 807 | soccer ball 808 | sock 809 | solar dish 810 | sombrero 811 | soup bowl 812 | space bar 813 | space heater 814 | space shuttle 815 | spatula 816 | speedboat 817 | spider web 818 | spindle 819 | sports car 820 | spotlight 821 | stage 822 | steam locomotive 823 | steel arch bridge 824 | steel drum 825 | stethoscope 826 | stole 827 | stone wall 828 | stopwatch 829 | stove 830 | strainer 831 | streetcar 832 | stretcher 833 | studio couch 834 | stupa 835 | submarine 836 | suit 837 | sundial 838 | sunglass 839 | sunglasses 840 | sunscreen 841 | suspension bridge 842 | swab 843 | sweatshirt 844 | swimming trunks 845 | swing 846 | switch 847 | syringe 848 | table lamp 849 | tank 850 | tape player 851 | teapot 852 | teddy 853 | television 854 | tennis ball 855 | thatch 856 | theater curtain 857 | thimble 858 | thresher 859 | throne 860 | tile roof 861 | toaster 862 | tobacco shop 863 | toilet seat 864 | torch 865 | totem pole 866 | tow truck 867 | toyshop 868 | tractor 869 | trailer truck 870 | tray 871 | trench coat 872 | tricycle 873 | trimaran 874 | tripod 875 | triumphal arch 876 | trolleybus 877 | trombone 878 | tub 879 | turnstile 880 | typewriter keyboard 881 | umbrella 882 | unicycle 883 | upright 884 | vacuum 885 | vase 886 | vault 887 | velvet 888 | vending machine 889 | vestment 890 | viaduct 891 | violin 892 | volleyball 893 | waffle iron 894 | wall clock 895 | wallet 896 | wardrobe 897 | warplane 898 | washbasin 899 | washer 900 | water bottle 901 | water jug 902 | water tower 903 | whiskey jug 904 | whistle 905 | wig 906 | window screen 907 | window shade 908 | Windsor tie 909 | wine bottle 910 | wing 911 | wok 912 | wooden spoon 913 | wool 914 | worm fence 915 | wreck 916 | yawl 917 | yurt 918 | web site 919 | comic book 920 | crossword puzzle 921 | street sign 922 | traffic light 923 | book jacket 924 | menu 925 | plate 926 | guacamole 927 | consomme 928 | hot pot 929 | trifle 930 | ice cream 931 | ice lolly 932 | French loaf 933 | bagel 934 | pretzel 935 | cheeseburger 936 | hotdog 937 | mashed potato 938 | head cabbage 939 | broccoli 940 | cauliflower 941 | zucchini 942 | spaghetti squash 943 | acorn squash 944 | butternut squash 945 | cucumber 946 | artichoke 947 | bell pepper 948 | cardoon 949 | mushroom 950 | Granny Smith 951 | strawberry 952 | orange 953 | lemon 954 | fig 955 | pineapple 956 | banana 957 | jackfruit 958 | custard apple 959 | pomegranate 960 | hay 961 | carbonara 962 | chocolate sauce 963 | dough 964 | meat loaf 965 | pizza 966 | potpie 967 | burrito 968 | red wine 969 | espresso 970 | cup 971 | eggnog 972 | alp 973 | bubble 974 | cliff 975 | coral reef 976 | geyser 977 | lakeside 978 | promontory 979 | sandbar 980 | seashore 981 | valley 982 | volcano 983 | ballplayer 984 | groom 985 | scuba diver 986 | rapeseed 987 | daisy 988 | yellow lady's slipper 989 | corn 990 | acorn 991 | hip 992 | buckeye 993 | coral fungus 994 | agaric 995 | gyromitra 996 | stinkhorn 997 | earthstar 998 | hen-of-the-woods 999 | bolete 1000 | ear 1001 | toilet tissue 1002 | -------------------------------------------------------------------------------- /image_classification/test_camera.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run classification on Camera, Press ESC to exit the program 3 | For Raspberry PI, please use `import tflite_runtime.interpreter as tflite` instead 4 | """ 5 | import cv2 6 | import numpy as np 7 | 8 | import tensorflow.lite as tflite 9 | # import tflite_runtime.interpreter as tflite 10 | 11 | from PIL import Image 12 | 13 | CAMERA_WIDTH = 640 14 | CAMERA_HEIGHT = 480 15 | 16 | 17 | def load_labels(label_path): 18 | r"""Returns a list of labels""" 19 | with open(label_path, 'r') as f: 20 | return [line.strip() for line in f.readlines()] 21 | 22 | 23 | def load_model(model_path): 24 | r"""Load TFLite model, returns a Interpreter instance.""" 25 | interpreter = tflite.Interpreter(model_path=model_path) 26 | interpreter.allocate_tensors() 27 | return interpreter 28 | 29 | 30 | def process_image(interpreter, image, input_index, k=3): 31 | r"""Process an image, Return top K result in a list of 2-Tuple(confidence_score, _id)""" 32 | input_data = np.expand_dims(image, axis=0) # expand to 4-dim 33 | 34 | # Process 35 | interpreter.set_tensor(input_index, input_data) 36 | interpreter.invoke() 37 | 38 | # Get outputs 39 | output_details = interpreter.get_output_details() 40 | output_data = interpreter.get_tensor(output_details[0]['index']) 41 | # print(output_data.shape) # (1, 1001) 42 | output_data = np.squeeze(output_data) 43 | 44 | # Get top K result 45 | top_k = output_data.argsort()[-k:][::-1] # Top_k index 46 | result = [] 47 | for _id in top_k: 48 | score = float(output_data[_id] / 255.0) 49 | result.append((_id, score)) 50 | 51 | return result 52 | 53 | 54 | def display_result(top_result, frame, labels): 55 | r"""Display top K result in top right corner""" 56 | font = cv2.FONT_HERSHEY_SIMPLEX 57 | size = 0.6 58 | color = (255, 0, 0) # Blue color 59 | thickness = 1 60 | 61 | for idx, (_id, score) in enumerate(top_result): 62 | # print('{} - {:0.4f}'.format(label, score)) 63 | x = 12 64 | y = 24 * idx + 24 65 | cv2.putText(frame, '{} - {:0.4f}'.format(labels[_id], score), 66 | (x, y), font, size, color, thickness) 67 | 68 | cv2.imshow('Image Classification', frame) 69 | 70 | 71 | if __name__ == "__main__": 72 | 73 | model_path = 'data/mobilenet_v1_1.0_224_quant.tflite' 74 | label_path = 'data/labels_mobilenet_quant_v1_224.txt' 75 | 76 | cap = cv2.VideoCapture(0) 77 | cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH) 78 | cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT) 79 | cap.set(cv2.CAP_PROP_FPS, 15) 80 | 81 | interpreter = load_model(model_path) 82 | labels = load_labels(label_path) 83 | 84 | input_details = interpreter.get_input_details() 85 | 86 | # Get Width and Height 87 | input_shape = input_details[0]['shape'] 88 | height = input_shape[1] 89 | width = input_shape[2] 90 | 91 | # Get input index 92 | input_index = input_details[0]['index'] 93 | 94 | # Process Stream 95 | while True: 96 | ret, frame = cap.read() 97 | 98 | image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 99 | image = image.resize((width, height)) 100 | 101 | top_result = process_image(interpreter, image, input_index) 102 | display_result(top_result, frame, labels) 103 | 104 | key = cv2.waitKey(1) 105 | if key == 27: # esc 106 | break 107 | 108 | cap.release() 109 | cv2.destroyAllWindows() 110 | -------------------------------------------------------------------------------- /image_classification/test_image.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run classification on images, Press ESC to exit the program 3 | For Raspberry PI, please use `import tflite_runtime.interpreter as tflite` instead 4 | """ 5 | 6 | import cv2 7 | import numpy as np 8 | 9 | import tensorflow.lite as tflite 10 | # import tflite_runtime.interpreter as tflite 11 | 12 | from PIL import Image 13 | 14 | 15 | def load_labels(label_path): 16 | r"""Returns a list of labels""" 17 | with open(label_path, 'r') as f: 18 | return [line.strip() for line in f.readlines()] 19 | 20 | 21 | def load_model(model_path): 22 | r"""Load TFLite model, returns a Interpreter instance.""" 23 | interpreter = tflite.Interpreter(model_path=model_path) 24 | interpreter.allocate_tensors() 25 | return interpreter 26 | 27 | 28 | def process_image(interpreter, image, input_index, k=3): 29 | r"""Process an image, Return top K result in a list of 2-Tuple(confidence_score, label)""" 30 | input_data = np.expand_dims(image, axis=0) # expand to 4-dim 31 | 32 | # Process 33 | interpreter.set_tensor(input_index, input_data) 34 | interpreter.invoke() 35 | 36 | # Get outputs 37 | output_details = interpreter.get_output_details() 38 | output_data = interpreter.get_tensor(output_details[0]['index']) 39 | print(output_data.shape) # (1, 1001) 40 | output_data = np.squeeze(output_data) 41 | 42 | # Get top K result 43 | top_k = output_data.argsort()[-k:][::-1] # Top_k index 44 | result = [] 45 | for i in top_k: 46 | score = float(output_data[i] / 255.0) 47 | result.append((i, score)) 48 | 49 | return result 50 | 51 | 52 | def display_result(top_result, frame, labels): 53 | r"""Display top K result in top right corner""" 54 | font = cv2.FONT_HERSHEY_SIMPLEX 55 | size = 0.6 56 | color = (255, 0, 0) # Blue color 57 | thickness = 1 58 | 59 | for idx, (i, score) in enumerate(top_result): 60 | # print('{} - {:0.4f}'.format(label, score)) 61 | x = 12 62 | y = 24 * idx + 24 63 | cv2.putText(frame, '{} - {:0.4f}'.format(labels[i], score), 64 | (x, y), font, size, color, thickness) 65 | 66 | cv2.imshow('Image Classification', frame) 67 | 68 | 69 | if __name__ == "__main__": 70 | 71 | model_path = 'data/mobilenet_v1_1.0_224_quant.tflite' 72 | label_path = 'data/labels_mobilenet_quant_v1_224.txt' 73 | image_path = 'data/cat.jpg' 74 | 75 | interpreter = load_model(model_path) 76 | labels = load_labels(label_path) 77 | 78 | input_details = interpreter.get_input_details() 79 | # Get Width and Height 80 | input_shape = input_details[0]['shape'] 81 | height = input_shape[1] 82 | width = input_shape[2] 83 | 84 | # Get input index 85 | input_index = input_details[0]['index'] 86 | 87 | frame = cv2.imread(image_path, cv2.IMREAD_COLOR) 88 | print(frame.shape) 89 | 90 | image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 91 | image = image.resize((width, height)) 92 | 93 | top_result = process_image(interpreter, image, input_index) 94 | display_result(top_result, frame, labels) 95 | 96 | key = cv2.waitKey(0) 97 | if key == 27: # esc 98 | cv2.destroyAllWindows() 99 | -------------------------------------------------------------------------------- /object_detection/README.md: -------------------------------------------------------------------------------- 1 | # Object Detection 2 | 3 | Use Tensorflow (Lite) + OpenCV to do object detection. 4 | 5 | ## Download the Model 6 | 7 | The tested model can be downloaded from: 8 | 9 | https://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip 10 | -------------------------------------------------------------------------------- /object_detection/data/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aiden-dai/ai-tflite-opencv/81971855d81c101f50b2abeaa683577ceca75bf5/object_detection/data/bus.jpg -------------------------------------------------------------------------------- /object_detection/data/coco_labels.txt: -------------------------------------------------------------------------------- 1 | 0 person 2 | 1 bicycle 3 | 2 car 4 | 3 motorcycle 5 | 4 airplane 6 | 5 bus 7 | 6 train 8 | 7 truck 9 | 8 boat 10 | 9 traffic light 11 | 10 fire hydrant 12 | 12 stop sign 13 | 13 parking meter 14 | 14 bench 15 | 15 bird 16 | 16 cat 17 | 17 dog 18 | 18 horse 19 | 19 sheep 20 | 20 cow 21 | 21 elephant 22 | 22 bear 23 | 23 zebra 24 | 24 giraffe 25 | 26 backpack 26 | 27 umbrella 27 | 30 handbag 28 | 31 tie 29 | 32 suitcase 30 | 33 frisbee 31 | 34 skis 32 | 35 snowboard 33 | 36 sports ball 34 | 37 kite 35 | 38 baseball bat 36 | 39 baseball glove 37 | 40 skateboard 38 | 41 surfboard 39 | 42 tennis racket 40 | 43 bottle 41 | 45 wine glass 42 | 46 cup 43 | 47 fork 44 | 48 knife 45 | 49 spoon 46 | 50 bowl 47 | 51 banana 48 | 52 apple 49 | 53 sandwich 50 | 54 orange 51 | 55 broccoli 52 | 56 carrot 53 | 57 hot dog 54 | 58 pizza 55 | 59 donut 56 | 60 cake 57 | 61 chair 58 | 62 couch 59 | 63 potted plant 60 | 64 bed 61 | 66 dining table 62 | 69 toilet 63 | 71 tv 64 | 72 laptop 65 | 73 mouse 66 | 74 remote 67 | 75 keyboard 68 | 76 cell phone 69 | 77 microwave 70 | 78 oven 71 | 79 toaster 72 | 80 sink 73 | 81 refrigerator 74 | 83 book 75 | 84 clock 76 | 85 vase 77 | 86 scissors 78 | 87 teddy bear 79 | 88 hair drier 80 | 89 toothbrush -------------------------------------------------------------------------------- /object_detection/test_camera.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run object detection on images, Press ESC to exit the program 3 | For Raspberry PI, please use `import tflite_runtime.interpreter as tflite` instead 4 | """ 5 | import re 6 | import cv2 7 | import numpy as np 8 | 9 | import tensorflow.lite as tflite 10 | # import tflite_runtime.interpreter as tflite 11 | 12 | from PIL import Image 13 | 14 | CAMERA_WIDTH = 640 15 | CAMERA_HEIGHT = 480 16 | 17 | 18 | def load_labels(label_path): 19 | r"""Returns a list of labels""" 20 | with open(label_path) as f: 21 | labels = {} 22 | for line in f.readlines(): 23 | m = re.match(r"(\d+)\s+(\w+)", line.strip()) 24 | labels[int(m.group(1))] = m.group(2) 25 | return labels 26 | 27 | 28 | def load_model(model_path): 29 | r"""Load TFLite model, returns a Interpreter instance.""" 30 | interpreter = tflite.Interpreter(model_path=model_path) 31 | interpreter.allocate_tensors() 32 | return interpreter 33 | 34 | 35 | def process_image(interpreter, image, input_index): 36 | r"""Process an image, Return a list of detected class ids and positions""" 37 | input_data = np.expand_dims(image, axis=0) # expand to 4-dim 38 | 39 | # Process 40 | interpreter.set_tensor(input_index, input_data) 41 | interpreter.invoke() 42 | 43 | # Get outputs 44 | output_details = interpreter.get_output_details() 45 | # print(output_details) 46 | # output_details[0] - position 47 | # output_details[1] - class id 48 | # output_details[2] - score 49 | # output_details[3] - count 50 | 51 | positions = np.squeeze(interpreter.get_tensor(output_details[0]['index'])) 52 | classes = np.squeeze(interpreter.get_tensor(output_details[1]['index'])) 53 | scores = np.squeeze(interpreter.get_tensor(output_details[2]['index'])) 54 | 55 | result = [] 56 | 57 | for idx, score in enumerate(scores): 58 | if score > 0.5: 59 | result.append({'pos': positions[idx], '_id': classes[idx]}) 60 | 61 | return result 62 | 63 | 64 | def display_result(result, frame, labels): 65 | r"""Display Detected Objects""" 66 | font = cv2.FONT_HERSHEY_SIMPLEX 67 | size = 0.6 68 | color = (255, 0, 0) # Blue color 69 | thickness = 1 70 | 71 | # position = [ymin, xmin, ymax, xmax] 72 | # x * CAMERA_WIDTH 73 | # y * CAMERA_HEIGHT 74 | for obj in result: 75 | pos = obj['pos'] 76 | _id = obj['_id'] 77 | 78 | x1 = int(pos[1] * CAMERA_WIDTH) 79 | x2 = int(pos[3] * CAMERA_WIDTH) 80 | y1 = int(pos[0] * CAMERA_HEIGHT) 81 | y2 = int(pos[2] * CAMERA_HEIGHT) 82 | 83 | cv2.putText(frame, labels[_id], (x1, y1), font, size, color, thickness) 84 | cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness) 85 | 86 | cv2.imshow('Object Detection', frame) 87 | 88 | 89 | if __name__ == "__main__": 90 | 91 | model_path = 'data/detect.tflite' 92 | label_path = 'data/coco_labels.txt' 93 | 94 | cap = cv2.VideoCapture(0) 95 | cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH) 96 | cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT) 97 | cap.set(cv2.CAP_PROP_FPS, 15) 98 | 99 | interpreter = load_model(model_path) 100 | labels = load_labels(label_path) 101 | 102 | input_details = interpreter.get_input_details() 103 | 104 | # Get Width and Height 105 | input_shape = input_details[0]['shape'] 106 | height = input_shape[1] 107 | width = input_shape[2] 108 | 109 | # Get input index 110 | input_index = input_details[0]['index'] 111 | 112 | # Process Stream 113 | while True: 114 | ret, frame = cap.read() 115 | 116 | image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 117 | image = image.resize((width, height)) 118 | 119 | top_result = process_image(interpreter, image, input_index) 120 | display_result(top_result, frame, labels) 121 | 122 | key = cv2.waitKey(1) 123 | if key == 27: # esc 124 | break 125 | 126 | cap.release() 127 | cv2.destroyAllWindows() 128 | -------------------------------------------------------------------------------- /object_detection/test_image.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run object detection on images, Press ESC to exit the program 3 | For Raspberry PI, please use `import tflite_runtime.interpreter as tflite` instead 4 | """ 5 | import re 6 | import cv2 7 | import numpy as np 8 | 9 | import tensorflow.lite as tflite 10 | # import tflite_runtime.interpreter as tflite 11 | 12 | from PIL import Image 13 | 14 | 15 | def load_labels(label_path): 16 | r"""Returns a list of labels""" 17 | with open(label_path) as f: 18 | labels = {} 19 | for line in f.readlines(): 20 | m = re.match(r"(\d+)\s+(\w+)", line.strip()) 21 | labels[int(m.group(1))] = m.group(2) 22 | return labels 23 | 24 | 25 | def load_model(model_path): 26 | r"""Load TFLite model, returns a Interpreter instance.""" 27 | interpreter = tflite.Interpreter(model_path=model_path) 28 | interpreter.allocate_tensors() 29 | return interpreter 30 | 31 | 32 | def process_image(interpreter, image, input_index): 33 | r"""Process an image, Return a list of detected class ids and positions""" 34 | input_data = np.expand_dims(image, axis=0) # expand to 4-dim 35 | 36 | # Process 37 | interpreter.set_tensor(input_index, input_data) 38 | interpreter.invoke() 39 | 40 | # Get outputs 41 | output_details = interpreter.get_output_details() 42 | # print(output_details) 43 | # output_details[0] - position 44 | # output_details[1] - class id 45 | # output_details[2] - score 46 | # output_details[3] - count 47 | 48 | positions = np.squeeze(interpreter.get_tensor(output_details[0]['index'])) 49 | classes = np.squeeze(interpreter.get_tensor(output_details[1]['index'])) 50 | scores = np.squeeze(interpreter.get_tensor(output_details[2]['index'])) 51 | 52 | result = [] 53 | 54 | for idx, score in enumerate(scores): 55 | if score > 0.5: 56 | result.append({'pos': positions[idx], '_id': classes[idx] }) 57 | 58 | return result 59 | 60 | def display_result(result, frame, labels): 61 | r"""Display Detected Objects""" 62 | font = cv2.FONT_HERSHEY_SIMPLEX 63 | size = 0.6 64 | color = (255, 0, 0) # Blue color 65 | thickness = 1 66 | 67 | # position = [ymin, xmin, ymax, xmax] 68 | # x * IMAGE_WIDTH 69 | # y * IMAGE_HEIGHT 70 | width = frame.shape[1] 71 | height = frame.shape[0] 72 | 73 | for obj in result: 74 | pos = obj['pos'] 75 | _id = obj['_id'] 76 | 77 | x1 = int(pos[1] * width) 78 | x2 = int(pos[3] * width) 79 | y1 = int(pos[0] * height) 80 | y2 = int(pos[2] * height) 81 | 82 | cv2.putText(frame, labels[_id], (x1, y1), font, size, color, thickness) 83 | cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness) 84 | 85 | cv2.imshow('Object Detection', frame) 86 | 87 | 88 | if __name__ == "__main__": 89 | 90 | model_path = 'data/detect.tflite' 91 | label_path = 'data/coco_labels.txt' 92 | image_path = 'data/bus.jpg' 93 | 94 | interpreter = load_model(model_path) 95 | labels = load_labels(label_path) 96 | 97 | input_details = interpreter.get_input_details() 98 | # Get Width and Height 99 | input_shape = input_details[0]['shape'] 100 | height = input_shape[1] 101 | width = input_shape[2] 102 | 103 | # Get input index 104 | input_index = input_details[0]['index'] 105 | 106 | frame = cv2.imread(image_path, cv2.IMREAD_COLOR) 107 | print(frame.shape) 108 | 109 | image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 110 | image = image.resize((width, height)) 111 | 112 | top_result = process_image(interpreter, image, input_index) 113 | display_result(top_result, frame, labels) 114 | 115 | key = cv2.waitKey(0) 116 | if key == 27: # esc 117 | cv2.destroyAllWindows() 118 | -------------------------------------------------------------------------------- /pose_detection/README.md: -------------------------------------------------------------------------------- 1 | # Pose Detection 2 | 3 | Use Tensorflow (Lite) + OpenCV to do pose detection. 4 | 5 | ## Download the Model 6 | 7 | The tested model can be downloaded from: 8 | 9 | https://storage.googleapis.com/download.tensorflow.org/models/tflite/posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite 10 | 11 | -------------------------------------------------------------------------------- /pose_detection/data/person.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aiden-dai/ai-tflite-opencv/81971855d81c101f50b2abeaa683577ceca75bf5/pose_detection/data/person.jpg -------------------------------------------------------------------------------- /pose_detection/test_camera.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run Pose detection with Camera, Press ESC to exit the program 3 | For Raspberry PI, please use `import tflite_runtime.interpreter as tflite` instead 4 | """ 5 | import re 6 | import cv2 7 | import numpy as np 8 | 9 | import tensorflow.lite as tflite 10 | # import tflite_runtime.interpreter as tflite 11 | 12 | from PIL import Image 13 | 14 | 15 | CAMERA_WIDTH = 640 16 | CAMERA_HEIGHT = 480 17 | 18 | 19 | class Part: 20 | r"""Enum of Detected Part IDs, for example, 0 is Nose""" 21 | NOSE = 0 22 | LEFT_EYE = 1 23 | RIGHT_EYE = 2 24 | LEFT_EAR = 3 25 | RIGHT_EAR = 4, 26 | LEFT_SHOULDER = 5 27 | RIGHT_SHOULDER = 6 28 | LEFT_ELBOW = 7 29 | RIGHT_ELBOW = 8 30 | LEFT_WRIST = 9 31 | RIGHT_WRIST = 10 32 | LEFT_HIP = 11 33 | RIGHT_HIP = 12 34 | LEFT_KNEE = 13 35 | RIGHT_KNEE = 14 36 | LEFT_ANKLE = 15 37 | RIGHT_ANKLE = 16 38 | 39 | 40 | def sigmoid(x): 41 | return 1.0 / (1.0 + 1.0 / np.exp(x)) 42 | 43 | 44 | def load_model(model_path): 45 | r"""Load TFLite model, returns a Interpreter instance.""" 46 | interpreter = tflite.Interpreter(model_path=model_path) 47 | interpreter.allocate_tensors() 48 | return interpreter 49 | 50 | 51 | def process_image(interpreter, image, input_index): 52 | r"""Process an image, Return a list of positions in a 4-Tuple (pos_x, pos_y, offset_x, offset_y)""" 53 | input_data = np.expand_dims(image, axis=0) # expand to 4-dim 54 | input_data = (np.float32(input_data) - 127.5) / 127.5 # float point 55 | 56 | # Process 57 | interpreter.set_tensor(input_index, input_data) 58 | interpreter.invoke() 59 | 60 | # Get outputs 61 | output_details = interpreter.get_output_details() 62 | # print(output_details) 63 | 64 | output_data = np.squeeze( 65 | interpreter.get_tensor(output_details[0]['index'])) 66 | offset_data = np.squeeze( 67 | interpreter.get_tensor(output_details[1]['index'])) 68 | 69 | points = [] 70 | 71 | total_row, total_col, total_points = output_data.shape 72 | 73 | # totally 17 points 74 | for k in range(total_points): 75 | max_score = output_data[0][0][k] 76 | max_row = 0 77 | max_col = 0 78 | for row in range(total_row): 79 | for col in range(total_col): 80 | if (output_data[row][col][k] > max_score): 81 | max_score = output_data[row][col][k] 82 | max_row = row 83 | max_col = col 84 | 85 | points.append((max_row, max_col)) 86 | # print(sigmoid(max_score)) 87 | 88 | positions = [] 89 | 90 | for idx, point in enumerate(points): 91 | pos_y, pos_x = point 92 | 93 | # y is row, x is column 94 | offset_x = offset_data[pos_y][pos_x][idx + 17] 95 | offset_y = offset_data[pos_y][pos_x][idx] 96 | 97 | positions.append((pos_x, pos_y, offset_x, offset_y)) 98 | # confidenceScores = sigmoid(output_data[pos_y][pos_x][idx]) 99 | # print('confidenceScores {}'.format(confidenceScores)) 100 | 101 | return positions 102 | 103 | 104 | def display_result(positions, frame): 105 | r"""Display Detected Points in circles""" 106 | size = 5 107 | color = (255, 0, 0) # Blue color 108 | thickness = 3 109 | 110 | for pos in positions: 111 | pos_x, pos_y, offset_x, offset_y = pos 112 | 113 | # Calculating the x and y coordinates 114 | x = int(pos_x / 8 * CAMERA_WIDTH + offset_x) 115 | y = int(pos_y / 8 * CAMERA_HEIGHT + offset_y) 116 | 117 | cv2.circle(frame, (x, y), size, color, thickness) 118 | 119 | cv2.imshow('Pose Detection', frame) 120 | 121 | 122 | if __name__ == "__main__": 123 | 124 | model_path = 'data/posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite' 125 | image_path = 'data/person.jpg' 126 | 127 | cap = cv2.VideoCapture(0) 128 | cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH) 129 | cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT) 130 | cap.set(cv2.CAP_PROP_FPS, 15) 131 | 132 | interpreter = load_model(model_path) 133 | input_details = interpreter.get_input_details() 134 | 135 | # Get Width and Height 136 | input_shape = input_details[0]['shape'] 137 | height = input_shape[1] 138 | width = input_shape[2] 139 | 140 | # Get input index 141 | input_index = input_details[0]['index'] 142 | 143 | # Process Stream 144 | while True: 145 | ret, frame = cap.read() 146 | 147 | image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 148 | image = image.resize((width, height)) 149 | 150 | positions = process_image(interpreter, image, input_index) 151 | display_result(positions, frame) 152 | 153 | key = cv2.waitKey(1) 154 | if key == 27: # esc 155 | break 156 | 157 | cap.release() 158 | cv2.destroyAllWindows() 159 | -------------------------------------------------------------------------------- /pose_detection/test_image.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run Pose detection on images, Press ESC to exit the program 3 | For Raspberry PI, please use `import tflite_runtime.interpreter as tflite` instead 4 | """ 5 | import re 6 | import cv2 7 | import numpy as np 8 | 9 | import tensorflow.lite as tflite 10 | # import tflite_runtime.interpreter as tflite 11 | 12 | from PIL import Image 13 | 14 | 15 | class Part: 16 | r"""Enum of Detected Part IDs, for example, 0 is Nose""" 17 | NOSE = 0 18 | LEFT_EYE = 1 19 | RIGHT_EYE = 2 20 | LEFT_EAR = 3 21 | RIGHT_EAR = 4, 22 | LEFT_SHOULDER = 5 23 | RIGHT_SHOULDER = 6 24 | LEFT_ELBOW = 7 25 | RIGHT_ELBOW = 8 26 | LEFT_WRIST = 9 27 | RIGHT_WRIST = 10 28 | LEFT_HIP = 11 29 | RIGHT_HIP = 12 30 | LEFT_KNEE = 13 31 | RIGHT_KNEE = 14 32 | LEFT_ANKLE = 15 33 | RIGHT_ANKLE = 16 34 | 35 | 36 | def sigmoid(x): 37 | return 1.0 / (1.0 + 1.0 / np.exp(x)) 38 | 39 | 40 | def load_model(model_path): 41 | r"""Load TFLite model, returns a Interpreter instance.""" 42 | interpreter = tflite.Interpreter(model_path=model_path) 43 | interpreter.allocate_tensors() 44 | return interpreter 45 | 46 | 47 | def process_image(interpreter, image, input_index): 48 | r"""Process an image, Return a list of positions in a 4-Tuple (pos_x, pos_y, offset_x, offset_y)""" 49 | input_data = np.expand_dims(image, axis=0) # expand to 4-dim 50 | input_data = (np.float32(input_data) - 127.5) / 127.5 # float point 51 | 52 | # Process 53 | interpreter.set_tensor(input_index, input_data) 54 | interpreter.invoke() 55 | 56 | # Get outputs 57 | output_details = interpreter.get_output_details() 58 | # print(output_details) 59 | 60 | output_data = np.squeeze( 61 | interpreter.get_tensor(output_details[0]['index'])) 62 | offset_data = np.squeeze( 63 | interpreter.get_tensor(output_details[1]['index'])) 64 | 65 | points = [] 66 | 67 | total_row, total_col, total_points = output_data.shape 68 | 69 | # totally 17 points 70 | for k in range(total_points): 71 | max_score = output_data[0][0][k] 72 | max_row = 0 73 | max_col = 0 74 | for row in range(total_row): 75 | for col in range(total_col): 76 | if (output_data[row][col][k] > max_score): 77 | max_score = output_data[row][col][k] 78 | max_row = row 79 | max_col = col 80 | 81 | points.append((max_row, max_col)) 82 | # print(sigmoid(max_score)) 83 | 84 | positions = [] 85 | 86 | for idx, point in enumerate(points): 87 | pos_y, pos_x = point 88 | 89 | # y is row, x is column 90 | offset_x = offset_data[pos_y][pos_x][idx + 17] 91 | offset_y = offset_data[pos_y][pos_x][idx] 92 | 93 | positions.append((pos_x, pos_y, offset_x, offset_y)) 94 | # confidenceScores = sigmoid(output_data[pos_y][pos_x][idx]) 95 | # print('confidenceScores {}'.format(confidenceScores)) 96 | 97 | return positions 98 | 99 | 100 | def display_result(positions, frame): 101 | r"""Display Detected Points in circles""" 102 | size = 5 103 | color = (255, 0, 0) # Blue color 104 | thickness = 3 105 | 106 | width = frame.shape[1] 107 | height = frame.shape[0] 108 | 109 | for pos in positions: 110 | pos_x, pos_y, offset_x, offset_y = pos 111 | 112 | # Calculating the x and y coordinates 113 | x = int(pos_x / 8 * width + offset_x) 114 | y = int(pos_y / 8 * height + offset_y) 115 | 116 | cv2.circle(frame, (x, y), size, color, thickness) 117 | 118 | cv2.imshow('Pose Detection', frame) 119 | 120 | 121 | if __name__ == "__main__": 122 | 123 | model_path = 'data/posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite' 124 | image_path = 'data/person.jpg' 125 | 126 | interpreter = load_model(model_path) 127 | 128 | input_details = interpreter.get_input_details() 129 | # Get Width and Height 130 | input_shape = input_details[0]['shape'] 131 | height = input_shape[1] 132 | width = input_shape[2] 133 | 134 | # Get input index 135 | input_index = input_details[0]['index'] 136 | 137 | frame = cv2.imread(image_path, cv2.IMREAD_COLOR) 138 | print(frame.shape) 139 | 140 | image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 141 | image = image.resize((width, height)) 142 | 143 | positions = process_image(interpreter, image, input_index) 144 | display_result(positions, frame) 145 | 146 | key = cv2.waitKey(0) 147 | if key == 27: # esc 148 | cv2.destroyAllWindows() 149 | -------------------------------------------------------------------------------- /test.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aiden-dai/ai-tflite-opencv/81971855d81c101f50b2abeaa683577ceca75bf5/test.PNG --------------------------------------------------------------------------------