From b821d1fb3c42e5cd1c11b6505050b90ddda945ae Mon Sep 17 00:00:00 2001 From: joosthel <146737859+joosthel@users.noreply.github.com> Date: Sat, 26 Jul 2025 22:38:15 +0200 Subject: [PATCH] add - initial setup for object detection --- LICENSE | 21 ++ README.md | 105 ++++++- USAGE.md | 64 +++++ __init__.py | 67 +++++ __pycache__/__init__.cpython-312.pyc | Bin 0 -> 3163 bytes install.bat | 34 +++ nodes/__init__.py | 1 + nodes/__pycache__/__init__.cpython-312.pyc | Bin 0 -> 189 bytes .../cv_aesthetic_overlay.cpython-312.pyc | Bin 0 -> 10305 bytes .../cv_blob_tracker.cpython-312.pyc | Bin 0 -> 7283 bytes .../cv_model_loader.cpython-312.pyc | Bin 0 -> 1984 bytes .../cv_object_detector.cpython-312.pyc | Bin 0 -> 3946 bytes nodes/cv_aesthetic_overlay.py | 261 ++++++++++++++++++ nodes/cv_blob_tracker.py | 187 +++++++++++++ nodes/cv_model_loader.py | 57 ++++ nodes/cv_object_detector.py | 95 +++++++ requirements.txt | 6 + utils/__init__.py | 1 + utils/opencv_helpers.py | 103 +++++++ utils/yolo_utils.py | 132 +++++++++ 20 files changed, 1133 insertions(+), 1 deletion(-) create mode 100644 LICENSE create mode 100644 USAGE.md create mode 100644 __init__.py create mode 100644 __pycache__/__init__.cpython-312.pyc create mode 100644 install.bat create mode 100644 nodes/__init__.py create mode 100644 nodes/__pycache__/__init__.cpython-312.pyc create mode 100644 nodes/__pycache__/cv_aesthetic_overlay.cpython-312.pyc create mode 100644 nodes/__pycache__/cv_blob_tracker.cpython-312.pyc create mode 100644 nodes/__pycache__/cv_model_loader.cpython-312.pyc create mode 100644 nodes/__pycache__/cv_object_detector.cpython-312.pyc create mode 100644 nodes/cv_aesthetic_overlay.py create mode 100644 nodes/cv_blob_tracker.py create mode 100644 nodes/cv_model_loader.py create mode 100644 nodes/cv_object_detector.py create mode 100644 requirements.txt create mode 100644 utils/__init__.py create mode 100644 utils/opencv_helpers.py create mode 100644 utils/yolo_utils.py diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..1cbcc7c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Joost Helfers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index 659945d..e3229ce 100644 --- a/README.md +++ b/README.md @@ -1 +1,104 @@ -# ComfyUI-CVOverlay \ No newline at end of file +# ComfyUI-OpenCV-Overlays + +This project integrates OpenCV and YOLOv8 into ComfyUI, providing custom nodes for object detection, blob tracking, and aesthetic video overlays. + +## Overview + +The `ComfyUI-CVOverlay` package provides four main custom nodes: + +- **CV Model Loader**: Loads YOLO models for object detection +- **CV Object Detector**: Performs real-time object detection using YOLOv8 +- **CV Blob Tracker**: Implements blob tracking algorithms using OpenCV +- **CV Aesthetic Overlay**: Applies customizable technical/surveillance-style overlays + +## Features + +✨ **Lightweight Integration**: Minimal dependencies, optimized for ComfyUI workflows +🎯 **Multiple Detection Models**: Support for YOLOv8n through YOLOv8x models +🎨 **Aesthetic Overlays**: Technical, surveillance, minimal, and cyberpunk styles +📹 **Video Processing**: Frame-by-frame processing for video workflows +🎛️ **Full Control**: Confidence thresholds, colors, opacity, line thickness + +## Installation + +### Method 1: Using ComfyUI-Manager (Recommended) +1. Install [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) if you haven't already +2. Open ComfyUI-Manager in ComfyUI +3. Go to "Install Custom Nodes" +4. Search for "ComfyUI-CVOverlay" and click Install +5. Restart ComfyUI + +### Method 2: Manual Installation +1. Navigate to your ComfyUI `custom_nodes` directory +2. Clone this repository: + ```bash + git clone https://github.com/joosthel/ComfyUI-CVOverlay.git + ``` +3. Dependencies will be automatically installed by ComfyUI-Manager on next startup +4. Restart ComfyUI + +**Note**: ComfyUI-Manager will automatically handle the installation of required dependencies (`opencv-python`, `ultralytics`, `torch`, etc.) when you restart ComfyUI. + +## Quick Start + +1. **Load Model**: Use `CV Model Loader` with `yolov8n.pt` (lightweight) +2. **Detect Objects**: Connect your image to `CV Object Detector` +3. **Apply Style**: Use `CV Aesthetic Overlay` to visualize detections +4. **Save Result**: Connect to ComfyUI's Save Image node + +See [USAGE.md](USAGE.md) for detailed examples and workflows. + +## Node Reference + +### CV Model Loader +- Loads YOLO models (n/s/m/l/x variants) +- Supports custom trained models +- Output: CV_MODEL + +### CV Object Detector +- Input: CV_MODEL, IMAGE +- Configurable confidence and IOU thresholds +- Output: IMAGE, CV_DETECTIONS + +### CV Blob Tracker +- Input: IMAGE, (optional) CV_TRACKS +- Multiple tracking algorithms +- Output: IMAGE, CV_TRACKS + +### CV Aesthetic Overlay +- Input: IMAGE, (optional) CV_DETECTIONS, CV_TRACKS +- Multiple overlay styles and color schemes +- Output: IMAGE + +## Project Structure + +``` +ComfyUI-CVOverlay/ +├── __init__.py # ComfyUI node registration +├── requirements.txt # Dependencies +├── USAGE.md # Usage examples and workflows +├── nodes/ # Custom nodes +│ ├── cv_model_loader.py +│ ├── cv_object_detector.py +│ ├── cv_blob_tracker.py +│ └── cv_aesthetic_overlay.py +└── utils/ # Helper functions + ├── opencv_helpers.py + └── yolo_utils.py +``` + +## Requirements + +- Python ≥ 3.8 +- ComfyUI +- OpenCV ≥ 4.8.0 +- Ultralytics (YOLOv8) ≥ 8.0.0 +- PyTorch ≥ 1.11.0 + +## Contributing + +Contributions welcome! Please submit pull requests or open issues for enhancements and bug fixes. + +## License + +MIT License - see LICENSE file for details. \ No newline at end of file diff --git a/USAGE.md b/USAGE.md new file mode 100644 index 0000000..e71dfdb --- /dev/null +++ b/USAGE.md @@ -0,0 +1,64 @@ +# Usage Examples + +## Installation & Setup + +1. Install via ComfyUI-Manager (search for "ComfyUI-CVOverlay") +2. Restart ComfyUI +3. Dependencies are installed automatically +4. Find nodes under "CV/" category in the node menu + +## Basic Object Detection Workflow + +1. **CV Model Loader** → Load a YOLO model (start with `yolov8n.pt` for lightweight processing) +2. **Load Image** (ComfyUI default node) → Your input image +3. **CV Object Detector** → Connect model and image, adjust confidence threshold +4. **CV Aesthetic Overlay** → Connect image and detections, choose style and colors +5. **Save Image** (ComfyUI default node) → Output the result + +## Basic Blob Tracking Workflow + +1. **Load Image** → Your input image/video frame +2. **CV Blob Tracker** → Process for moving objects +3. **CV Aesthetic Overlay** → Visualize tracked objects +4. **Save Image** → Output the result + +## Node Parameters Quick Reference + +### CV Model Loader +- `model_name`: Choose from yolov8n.pt (fastest) to yolov8x.pt (most accurate) +- `custom_model_path`: Use your own trained model (optional) + +### CV Object Detector +- `confidence`: 0.5 = balanced, 0.3 = more detections, 0.7 = fewer but more confident +- `iou_threshold`: 0.45 = default overlap filtering + +### CV Aesthetic Overlay +- `overlay_style`: + - `minimal` = Just corner markers + - `technical` = Full boxes with corners + - `surveillance` = Thick boxes with crosshairs + - `cyberpunk` = Glitched fragmented boxes +- `color_scheme`: green (classic), blue, red, white, orange, purple +- `opacity`: 0.8 = semi-transparent, 1.0 = fully opaque + +### CV Blob Tracker +- `tracking_method`: + - `background_subtraction` = Good for stationary camera + - `optical_flow` = Good for moving camera + - `contour_tracking` = Simple shape-based tracking +- `min_area`: Minimum size of objects to track (500 = medium sized objects) + +## Tips + +1. **Start Small**: Use `yolov8n.pt` model first, it's fast and good for testing +2. **Adjust Confidence**: Lower confidence (0.3-0.4) for more detections, higher (0.6-0.8) for precision +3. **Style Mixing**: Try different overlay styles with different color schemes for unique looks +4. **Video Processing**: For video, connect multiple frames through the same workflow +5. **Performance**: Disable unnecessary overlays (grid, crosshairs) for faster processing + +## Troubleshooting + +- **No detections**: Lower confidence threshold or try different YOLO model +- **Too many false positives**: Raise confidence threshold or IOU threshold +- **Tracking not working**: Try different tracking method or adjust min_area +- **Overlay too bright**: Reduce opacity or try different color scheme diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ab29714 --- /dev/null +++ b/__init__.py @@ -0,0 +1,67 @@ +""" +ComfyUI OpenCV Overlays - Custom nodes for computer vision effects +""" + +import os +import sys +import importlib.util + +# Add the current directory to the Python path +current_dir = os.path.dirname(os.path.abspath(__file__)) +nodes_dir = os.path.join(current_dir, "nodes") + +if current_dir not in sys.path: + sys.path.insert(0, current_dir) +if nodes_dir not in sys.path: + sys.path.insert(0, nodes_dir) + +# Import nodes with error handling using direct file imports +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +def import_node_from_file(file_path, class_name): + """Import a class from a specific file""" + try: + spec = importlib.util.spec_from_file_location("temp_module", file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return getattr(module, class_name) + except Exception as e: + print(f"Failed to import {class_name} from {file_path}: {e}") + return None + +# Import CV_ModelLoader +cv_model_loader_path = os.path.join(nodes_dir, "cv_model_loader.py") +if os.path.exists(cv_model_loader_path): + CV_ModelLoader = import_node_from_file(cv_model_loader_path, "CV_ModelLoader") + if CV_ModelLoader: + NODE_CLASS_MAPPINGS["CV_ModelLoader"] = CV_ModelLoader + NODE_DISPLAY_NAME_MAPPINGS["CV_ModelLoader"] = "CV Model Loader" + +# Import CV_ObjectDetector +cv_object_detector_path = os.path.join(nodes_dir, "cv_object_detector.py") +if os.path.exists(cv_object_detector_path): + CV_ObjectDetector = import_node_from_file(cv_object_detector_path, "CV_ObjectDetector") + if CV_ObjectDetector: + NODE_CLASS_MAPPINGS["CV_ObjectDetector"] = CV_ObjectDetector + NODE_DISPLAY_NAME_MAPPINGS["CV_ObjectDetector"] = "CV Object Detector" + +# Import CV_BlobTracker +cv_blob_tracker_path = os.path.join(nodes_dir, "cv_blob_tracker.py") +if os.path.exists(cv_blob_tracker_path): + CV_BlobTracker = import_node_from_file(cv_blob_tracker_path, "CV_BlobTracker") + if CV_BlobTracker: + NODE_CLASS_MAPPINGS["CV_BlobTracker"] = CV_BlobTracker + NODE_DISPLAY_NAME_MAPPINGS["CV_BlobTracker"] = "CV Blob Tracker" + +# Import CV_AestheticOverlay +cv_aesthetic_overlay_path = os.path.join(nodes_dir, "cv_aesthetic_overlay.py") +if os.path.exists(cv_aesthetic_overlay_path): + CV_AestheticOverlay = import_node_from_file(cv_aesthetic_overlay_path, "CV_AestheticOverlay") + if CV_AestheticOverlay: + NODE_CLASS_MAPPINGS["CV_AestheticOverlay"] = CV_AestheticOverlay + NODE_DISPLAY_NAME_MAPPINGS["CV_AestheticOverlay"] = "CV Aesthetic Overlay" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +print(f"ComfyUI-CVOverlay loaded {len(NODE_CLASS_MAPPINGS)} nodes successfully") \ No newline at end of file diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74424af6fac3fa7999926f0bdfb3d1fe17c17d06 GIT binary patch literal 3163 zcmb_ePf#0G7=QaG$tHn>kc5Ol>9*8Xuz#pp(YAw=0s=}y6=kNeGi%s=z{Vw;*z^~8C z$$C24z^@*!=pkNFp~;D_4Yna(ZHQMB=GBaNVW|OreO>_vFV?`X990-vwD0@YK0@)KS4*3S<`I*9dW9&p8 z<_4$OiAAVLnS#omVFwpfO`d0SvH(?fMpjr}p3g66P+=EERg`lqoSA{Vruu$0+c0Ig z@J_yVL~(0+(MX~HlN0FRZvZdQu-Y85Pw%T~vo6Q%2m(<8cHp}B5W>`pc`lb+u;OIx^B^0$X*L%i9iXbN&t;?iT)gX z;kJ1Y;5raox9Hn+!DY=$#1gnkT?b1f@1Exja88}x!EN=D3t<4kDI z#N#BH&uFNoMvH}MnaoXwTaq&ZR1817iCiqR;hBY6M$;6-JA9LedF<-SD`HOTaT#6= zna^l*h7ZTU0s&RC+N}BfIW=_gMy(O;nOhQbg1n@1W->C<60{9Bc-)s9 z?lxO_W^k(3xi>gY%!wMuoy`}Fs1-Nfdl0Q+p3^q;S3gF>GT5g*uE<_6zIuD*_T%7% za`3{#`A63)!HIkHUMR8AwchnO)K?DmJ&aUBuiSH2nWk0$ivKZlw9FjccJDB!_JVC| zi}#C_V6V>f{;}te{0>Nj$TaIM$F_U7J0DCv9D8`_o9RcPN0c5I+F^$Gg5kgRU1-yv zY8-L+B>7zOdGPJ?-YY@s+w+ktKI%K43(3V+e$h-hlu(7tDwsh92B)~IrfwXoia~U2 z@tUnEqR2^}&Z4u)2~b0mjQv8$Zvv}_YT9C}dE&8(JtCp{mX-NK&Dv)9`Sd zAoCn&kg{q}YC$yu&dE4y9t_3&L^5dHyB^LKGsF)lEb;FyYs;P?EkcOGZL#4yJl$Z_hbZIP; z*+52ny#iQv^Sj54JR?aQr}+Pm1p1 z&Vb%}a)&y#8|~F!cwJAvU5UQ4oZ9t7t9=7H)1h~sE;DDU{i8Z_TPtfy@?4qet6ofD z$XCitKZ5-KNgg@6HL^LfHNH9ipk06YYNa!!gZ5`gzOX}`-c6j-`^NOQKBy!zIvv~f zB&rustSIjDW#+}|<##a2SIf+6Klxex_<-(vecAIfaOupkFT3@VuT_p;)Q?@#FQs(% M#1G(_(uX?YZ`C09YybcN literal 0 HcmV?d00001 diff --git a/install.bat b/install.bat new file mode 100644 index 0000000..256b3f5 --- /dev/null +++ b/install.bat @@ -0,0 +1,34 @@ +@echo off +echo Installing ComfyUI OpenCV Overlays dependencies... +echo. + +REM Navigate to the custom node directory +cd /d "%~dp0" + +REM Find ComfyUI's Python executable +set "COMFYUI_PYTHON=..\..\..\python_embeded\python.exe" +if not exist "%COMFYUI_PYTHON%" ( + set "COMFYUI_PYTHON=..\..\..\python\python.exe" +) +if not exist "%COMFYUI_PYTHON%" ( + set "COMFYUI_PYTHON=python" + echo Using system Python... +) else ( + echo Using ComfyUI's Python: %COMFYUI_PYTHON% +) + +echo. +echo Installing requirements... +"%COMFYUI_PYTHON%" -m pip install --upgrade pip +"%COMFYUI_PYTHON%" -m pip install -r requirements.txt + +echo. +echo Installing package in development mode... +"%COMFYUI_PYTHON%" -m pip install -e . + +echo. +echo Installation complete! +echo. +echo You can now restart ComfyUI to see the new CV nodes in the node menu. +echo The nodes will appear under the "CV/" category. +pause diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..f27a938 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1 @@ +# Nodes package diff --git a/nodes/__pycache__/__init__.cpython-312.pyc b/nodes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2abf7e5172119cf5c8f50fc96d2e80cf177b9bd3 GIT binary patch literal 189 zcmX@j%ge<81n-SoGX#P3V-N=&d}aZPOlPQM&}8&m$xy@u z38Ul7GxJjN%ZuX+@{39mlX6mF5;4i8#U=T<@p<_vsl_lsUFR_WvecrS#L5_mczk?j oUS>&ryk0@&Ee@O9{FKt1RJ$Tppk<6eTnu7-WM*V!EMf+-0J`ln#Q*>R literal 0 HcmV?d00001 diff --git a/nodes/__pycache__/cv_aesthetic_overlay.cpython-312.pyc b/nodes/__pycache__/cv_aesthetic_overlay.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b4f2d0fe6f4e97013ebd9048f6c4837e0748c2f GIT binary patch literal 10305 zcmbtaYit`=cAgo&Um`{7K~j`NJ#349NUmk z*%mX(2HR{)1L=n97B$s4&3MsbX*k6yzyfu#0oMD^K!F7rIcx63K(*}xg@0AH6Lhir zqdn&ihm>u4H$lhNy?5?C=RVIp=R0TkZw`l*f`>c!{h42Gr>OtL7wyrP3XlF03RftW znxrfgOS4Rzo}@|6Ofn?bP3j@n#SL@DN#mSp(llqDG}9ETpR_R4AjKNrrC8GqM(uC1 zW|8TjIBO@x*G<}31JZDI&US+?wL8vGEHi0)!}eaQnqQUFsXI|#K~Z6E;0>l7h~c~kmZ0|j3pC6HX=r>_i=)wA)R6v zMk2*RvZH+DVySp6nGh$AMb&kvX=%~aqwm1`6^f%^>nPl|NgbzWnRlrhCF)57tD7{! zmYHBn%&c+J0-Nn*&F|8aR@U+^1zTl>)Xv%UHg*@?h*M&^n5HaO;#Q)9SxbP4N@KaJxW0WVnI*Cas3F6pVi0cDTXz$ zM%Kid;m>ka5ARmi#@bm2Tcf@MwUbc02z4#+xvP{sPt?_2wXpSU1KW6&VZB$GaVq2+ zy-vePoah@I2)Ilk|yX?(al9JDSAQV<`u)Sk&}HB(|^B$&;7&y1vKQDR^2nV85a#v~s}fczNe7x?)&$NPa>G4-E3IWo{UI&myy zQ4B;16-Viu!bDn?ns_Y1g~gdz^n8L71jUq`k3?hQqGFFG<4Ha&L}$1;PH_n{$&2A= zGRd>CL`38S#aXIJOvP9Zq+hX+ifKN^YHfT{5N0AVUI@{O1ElEaz{J3jiQ$u@<4|FN ziixqlBd?5y7{yw?6@p?UVj@7TndjbIh~Xm5$$7jok@$5ASas33AmF|Re?N2dz}eF6 z)bjAfSb|Mn6vFdKUW}ZJb6UgM=z<_7=fVlN)`C{K>&R;*nLex5M=ylSG951MRrma& zVjCVkb$TK^apu&(c&fF^o#@^~P!E$6ZbI@7wPB|Gol5)b#k%I~TN^rNl6G%UOuJ{p zNOgvat-)ebpxD%03^YHm7w0b%pI+$-h^jOP-9p^ZO;d!d?tB-l#~{(H-UszdG|=O1eB#mQXKSBmfLE=JHKYs_HS2krLafv3 zW*f9pCD&HCWEXJ?N)M_PGueC0DCGQkHPIXY?A*~gVU0}(#NrTjihqX4z zr*VEjeaL9#@6oJ#$@jKz6Vmgff{0rdgJlh`!L%|;OEsK`2?7Br+&q^+(*}s|Krl7X z9Xu81A_5nTB?K`Nj|WTEX^u;XKrkPhFSjItGkRgy{GvDmI8RLS(V5^WK#9qV_ff=o z0u25aN{W7Lu)hzG0oFW`gmh{JZX_v{D(?RsR$whuV0IrxgvT8%Rgl`pDEpe?`6RE{PYjQb504Il-9I(1SYxoOuyF_#3}j0b;{=xw zfKbQA=D?`WFDQn@!W`IIV+6p-Jf|=TfJY0lg!lrFs!Xv=@yWRe&qo#&U4pv^04kQ^ zc-&G&Pav}nXccRi1^66Rk%(F@!wbcxRuGJ#xYU0jGBtAbA z0o*e;pA-;g8Pw5DQ}JX(+|#RAaYkWv)|Q2Y@a6)?r2q>JT#9l8vM4%1KMd+l7mIWr+T*)0f@ykM%zeTnpTBb{9R~D=#j;_``#8+h?ARVxUcKTN%m>6#cCq zg;qiz?6~XSQSk3r^Y2>K75sbd`VJO+2S4rq`Ov=(t@#Er#(O~5v(l5>y5`@#Y+CoX zW=vlism8#S*m6uZ=1f1aU$?Ku*BXwb2k&_TS6*6vN#3^R*`7X%pl(aiyQSFJS@d-k z{oCOG1Ba=hCVgzfO*MDOhYEq`(kJfKHD*n6$BJXEZfkm=*t$<{&bRJcH56L+B@JOzatN<)o)K9gYEM*ePmuS%MEMZuDjlzg16_V z%)eMav95U!-sB42BkAFLfwsGWorS>8+?ln&{=0#`LZI(f+gjjg`owx&Ti(-=w|D&R zw@&XviiQIxlJRGsMsKzKZ0hE>@~y|#+=F@h;2$=0*!p|H3zGD~kZ!1zy4`9S*6VI} zH4Yut-9Bu9auk4C`6&Za7q|{rD2X`>;8bhX39cs{E#l;Jv0=8bi6ULSs+b z4&v%=%)FW1A_sm|+nw$!+FkDtUmniV+1;57GA-}RzLjs-vAX3`;nwb3vHYuJzlg6L z9#0Q{Zl8dQtnzP_cmArj=ShBp{AM8ZV8zf9LE`6tlTSc08e)hH z@jSddB{sZ(bRs4wF8l%{iaAR93|}gj7L}Uv0AWN1Rb^)!m{Xn4mYs4JLIx;>4qVY2845dr!VR169bOB3(SESb;hL19^ztV^~fyQE_o;uCvW-I4=5q_QXh zf9a8EP+)pcDUKT@6-QHvp980!aCSdKKrE6~vPt%7J@Rf=Zhw3f@P8_!ya%J0$S8@g zGYV@)A$)8si)6z()`FJ-W@N9tEjgEJSSxt8)m3+m(*^O72}*D=`-+;APlaA~~Mwy+|I` z{XyAhgUGyx@YOwonqfWI+orMHd>30+EovjVH7fA^tHx<0i&P7{TL%iOzH%Om4frqA zSW-G=kgGX)5H>xRyu{;<@EE=!p<0MDnvx45F zLU$?jPK6FBR`3=|LC}=0yLSpI;t@dzCn9rPN*9!Z5JF1=$7sl`I7%Kzk_Q_pAl#0E zOEe7uThx4-LoywYiP4!bP8RJb_(SvPiz@orB*c#OB6mp))o-$`Q%GbW0o{tORBV-) zEsRtS*Ya;8HFIgV!YsnyrC!K;6jl3EDN!nol@#Y*z|2rmkX z9@oyl2c!bdv|^|?l~kMkDm14$o~?Jwr}Yom*?)sN0=4V(W*Z95U|L`FHD9htGwJ>e zlj&c#xiZ6_*_zk=f%Ge2vhUhk3ig()NFZ#U>?2Pzj(u)_85t9wIXbH|-j4iY7c!qJws*?>wQae`wU*WPRsNH0 zHzS|4WMekz0imjpaCsFW7z^Y8o&ZLUZ-(H#2#k3 zD>qqa-(O<#=8hNI_C3u6`X89U+SGWomEN_t73^(tXHK8vR?pnA9|kYN-nbIFSMSZb zv*%VC<@UAuEdbelP3hymar%oc|6NyG!PO@B$>(76YID80&Rj(9zT?{c8`%1@xiesT zij6@zP-yJV=qvOSxeodG9oKW~&Zg{Q!P!-53FSr#?JwSO9r%sYR|HQgwVaZ-uQl#M z?z$_MWlOdw4+ zu}nn~vihqGAatrK4npGwhzBwM=I|kFAfxocC?xDjyaX_sfWM4*FP3c2FdJxa5Lbcm zBYDXHk$2!D=1~7}8O_twqIuCAvZU$>LLzQBhMh{DJ-Ma)_koc|4+roOURzQD?mmG& z1Vcds>K6XHkl<%T-83iic-IsY7mvs01ukUYq$x1S!~ZoV6-@O5e0?8TYNjIRxp4U= zZq@_ypm{_URbjM$)AbIj_j?R#tA7{S6mj__5D7nm1n#!2Hgo#&VGv+nOLl6-QEch> zD7F&I*;k`$EeGIpL`zM26i|_)_O3lxum|Nj`R&z(n;gREH}c<}x?`VSw|O!P1zTXf zZL56gS}W*)^iXD|=xdi>U1?Tt;mGBYY?nM)@b|2GZyqn~2H9_uj~9I1tLB>r3eOF| zP4o}a>&@-C0gyz8?<Jf;yWG8vVuw&a3>b^bc* zM7LGjp>M*etrPjd4f#gvtjzWqoQq6Im+C0F#AvL85Li18HaiFr-E+x>1bhdbhQcM) zf4NRS?YfW)Te)U`3v$GxPr)x^@;8|L5R$51%3nc4^)T9F0Ad~a539CvvxTfUA^lk< z^FsEVe6Y~cy~^a9_vGEr-`soKlD6D2j}^^!;;Cpv?ebfN_TDFmHuu$pfu(bKPbk;F z%H&@7fAz&=-za#yzE)pVC;!do^^5^OH;+B-hoh1swujw(*&$kl*BxyJ2pvR&RH0Wr za&pDhc#i?1F{L*ASonB5w5U`h$2SmPBZ2TF=7jwaVTKmyr_}a%KQOwfab3?I5!49U$|`fkdYthzDnbjyso?V`Sj_4w(fVYZ;XSXOBC&$) z>UGn8Qc?OnjS7Z2_OV-cAJ0Y&B9Qn2FT;Dtw#m{xar>%VJjE<<)8y-oV$w?J8*H*| zs%g6nm47T`P=V`T0;SMf*xMnB)t((&*|NMS zw=GKrcSm*hR{;q3UNfwm$u(mq0;=c z(QD;?K)wUle7U}BfkI$cg}ogX*Slhn>E+aN&t~?GT1%ZZZGO~1d4uV}5;)FG%LA)E z422zpP?$GJu3(569Szkg<}h)d!eRK96NWD`3vmw0_Hg)3_+=HkRq$pChuLHl3O3?x z&2jLpjpg^i0l;Udfr-;&qv~%p%5S56CkDn9^Rd&TkECuzw3Jn`krEQTVC(~AG@HV4?UzX{oRI*GJ76U4DE-o88rUd(~x=dhx@)i@`nur aRQ_JT-Fw%&jry^nkJ5c?J*fK`7XA-*Rtk3j literal 0 HcmV?d00001 diff --git a/nodes/__pycache__/cv_blob_tracker.cpython-312.pyc b/nodes/__pycache__/cv_blob_tracker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1d9e5eae4a2fb8424d1d297acd1d2f62b31aa41 GIT binary patch literal 7283 zcmbVReM}o^dY|#k_}ds_V!&V+10it|lFb*J4Y?2pOu~j37T8?EWp8Rc1DF^aX2uYF z<4x`^D-ktSjrI>hlT(~-rIIFpaOz3Ts*x^LSJ#vJ501BoZ#I>XHdUMdNVuw}U7gzJ zo$(k$*sI!C^7DQ?&->2%`Fnov>wmV{%oqf&^^?&<)fo0Q%4mViC)PiP#5~4ggO~|p zaaI?>2XQ6V4eFGb7$hMkBGkBkP>*9QIcU&fXE9cP6JrgJbZR?;))`#`#u*zizI@Qc zQpkogbEZdlUT?XKvARLi4@|e~)OZ@7#;y^Uv1#0EOfcs5i}aaDY=|D|$oVU~;0!+da@6X>{r z?0P1+X!}L;E7XgC7niB7@!$wI&>2#@rONymu=QTZ&tn`0Gk~Gd4-%~ICN@ZN6ieL1 z9_8x{>d_p)7>%qR#%dTevo_Xr6Cbp&X0W!h7LYd93bJIz=q-^;+AlI_aQU$%PUOd; zbpBv|DB2*l={_BdjKuh`I65wb^7E;51@)VN-}(U%XbmNdQ4~lhxq|fw(rM$wX0V5N zz)Q&H37(r`lrPFwQJr~4oSEQcV?HOqoRI<-85U4e^f#f76Ib#pqQ%T~ILgMR1!f|~ zi@~7?rzx(4;({0(XQHr_f>zkvez9kY<0HYDD{6UYib3|Vn#MHOgI`A875J& zdFBpiik;9Ip%TOmEa$47C2!kvHS{d`rLihk>6#_~%(f0XY|~_!pg@_dGmwDE!bf03 z{v~W|mJH*=I+kMftbsMc&vcJe@&?w-T3Bm7Z@WjbC9FM+v!(ZR{g}5*-A)&LXZtpga#vqV^P@}n-IgHV1yZt#HQuaP%J9O;=DRo*q^e2<%Wath{%^g zuVkVtFcAJ7%=D|nkPVma#xO6*mn*9Aa8xG7gV$g=L~i00I@z!O8wm=8mz1e6+S#&Q znMY`v8Rx{&7%Lm0A`|4fpllLHc}^IOMOfL67Csz{3u;f$Dw|-c2l{;N=li`n*~oK~ z@i5Opbu@rjG#Gh^fe{gk2DlGQ6C+$Gj&x-nWyZN!pZtkrU3Xpy%BZ9bERs7}%+4I^06DCt;((ssFojqTb} z0kV=oA|&!rer=P+=sQrusMK&k4LFrs1Js}jUou83!QND`_inX*`wNm`Y~R~f+WG)= z7iRzev6Bc1Pnx3LB0@kVj3IxqB}jU0-p3AYQ9)Mg32RPRV6T}a^B>)t>dKXY)seT` zVoB{wsGX!qGD*f^K-30IJgQj(ji^alicXhBd{@&+7OnnW>}NVH{S%zECM_RZifDx` zp|S~DB?;juEsk<@D9-a-R8(m~!g(PqC`5tffGMzCG!*896LjL#0lGKB1qF@{M+Gq$ ziO_jO9EZAKFcF?8Xkt(^G}SyY19%JHi1DG(R|pGu!K4f6%R+n<1!WG`QVI9M7XF-38O(E14L-*uxt!XiS4lnv{KUE)7{g@^qoC(_-vo= zvP?7EJHHJt=VXfA;6`~F^Oi#-L$Qe&9zmpR9_A5yP&!J1{|t(Cjl2hX;UReN8`!_ucfHtAyX<-Qsdee&H4N*-J9H4y ze%(5TIXY3=29n%uI)lHoSKZn5S3Cc@=0VMu_5(Rb<@|ed?|s^u_B?kqkp8Nb3LZ0D-I zcDeqzA%A*PYJp@GVwlz`;9^)?{=1T zHeey^y-ZCHn7C~6Ju z{4JqKi%8X$pm|E=TgIebGIn4e&A@j}MPEfMY)g;|*F)q9B;C~xI6V06+M-$rF1Lh@ zT3297r{v0qj-p zfEwcwMgDj7duweENtqlPid~cSA@DaGFOv#`Hv?P_!zCA>9q)Z@vh* zU4c{pdbP|kt6#O)|Ge#1+p?wUrNcSjI@fwfynlV+`VyIS9GmTUQCWT8wqQ%Mi(N~? z)Ap5<7cwU=WGe%=x@LW|Q-EUCwJWZ@8Q0#$mR}rMInt3i((y%TrSqdq=SNxB4{x2H zJ(mi8UEPo}Ak4i!cRf9wtvHw>bB>BTdp>P@;b~m)9LjhOE$z;FPRzZZ@}(y4cfp2kk55vpPNJsJ2{yuMD{d7UML_}!bXOTE+y>yx)S@%P|3*aD zSS5^Yhj`@1?P@!;drW92GD(UEr`49WL9*>fBjFM}auqNW`i?LYZVGzIz>!<#yRx@u<6e^s?7MdcjTcjC_-jk)r=mGY)cdD9}bIQr1Oy!-gS zl)tw|VJ`1$BUW1dqL#itx-gpl(bM{0?)`l4m$m+s1<1(FKO@Mp5M?%=FR~+#o=M*e(OT( z-S5qHAh61*aP*>~DLwgF%i|-9laE@K>ksFeyZ~+K57PXDfzPTQw=dyK^*<*cb^Stk za_C9l=STl`=8J}B!C&kCf&EACGTr^!Kq71+0h=hyb9Ak=2%C!8Q-Obd0kfp?FA;4h zC5+`gvX)0NoNC0A3Psy-vBH^$`bjFOFKTf+9MTuHwhikl(r&}DHt{4wl9CL0o@6Zg z>UNlA)Xu4LTi)PF8>8Qbl}UwK3Qq!TLQDymjJ1dtz0I4*?eGg~Sx`-dC*}DP@T3yS z4m`n`6AQFh4Y#ec(MOhAsVI&kDfjI5vV^7oS z!S{hCL2o_y`B4@wz@LGY@E;&KQ)1$fk*vFA&YZ$ihf={iWx48_JCkYO!qmcvlyOsLwbVIpo3rKI zwP`$kC_VY$*rI1K_;A;9%{#fe-HVmWwTI#5{YCvQph%-n${t2CO|4J*v-Lh89CbT@ zaFjWKa8%O8go8K0h`o07%pZlnY8XZ~Gt795jYm-0$}p4hU?i_GFbo?DF$|BMLV=h@ zf`XtvZ1cNOi52MVhkXI{1-g(6_%8VSWn<@uf%bu}o`7s@_YL^Z_Vis=?(aMzj{IpP z&@%Dp1yEr^>e3-1!Xtt19>RGLH?iN69dAg2ss?n_m%Y|iQ08?4Qfnx+b{H$)d!sw& zXaZrcxp7Vr>*Rjw4F2W{?xLJ;j@MGo_0xDI-f&gU%vhk!ko2U3spbUY8fX)zHcyau- zj=;$^s}Xm7MV5R;mVQM#|COxDxjoCI>%XjF9qxRCA^H0?3ubh@!F0GQ=dK6kZynCm hLE$$7+Wa@m8?eu)lf7%Q6rR~1|E!%liPMy%q}k9=U6`q2I-!v6MR4TLN$$5tK&+6UoMT9IhcLF9Aok<)V^TFd|A` zmCz_6flWk$cct3@d2LlXj%c7Cv5tI@s6vN^Xz(ts%foYsNcrFg!H;`u`!ZfeEAkv# z#)(iRqoro`afeV_A9D(XGCd0iecoYu$|=sT&SW&f=la~l*o0mbirkkv7e1t>s}t(N zW;rG8FL7Tsfjvd2PLG(i`)joBEZQ4{^L0c~9z(zoUS6gOkv7rYx^`YAay|f|icmr% zkPnkE32tHtGnCN0L#bKA_qOJ2<<$}4l5!{!$$FZf1*23b(uBINtva@|G*}ua zyOlcTO}N;E?Iyee*b5RmUs$wV76K={Hgl1?PSL2DFBjYeFOZv_%4A19G?I{5SSW>y zpl{h$i6$g3#OS3(i&5f*oU$lF!M=l_0xFFW#Owbbmy5@mK_^?Z$QF4aA>{eBrJuX>S8c5-Ign4X(V=PKRLmuR3# zb{JgZX&@V@8bWP7KlRd`j-DH9RasiV?Nuc8#;dZ@b*LIeZT+7u?DY2S^qtsM<<~(b zp8%QKb!fLG+NxDiG_37`ArWmzTsSeD9s{iU$@4VfTg}Z%I5m69Zy9cif;UzWfZqUL z0=SOc1|RXbd!R?K+Za=H6B^{U+!(7XFr;e|lh&jXUc>9MEA07e?LcQ;SySe54jo6X zD6qd^OkZ4An{JTZNJF-!Hmu)7x1|QYf{C)OURImNR;|QD-#5)l{u`rnAH8Y$}t^8OhmXW-K{0miD4kizU}8(llcZ^P=OKTn;+e z0PP$TdG*5S6_fhi>dD+?U`pZCo@XL~Ogxv30TUCzIt2q$F6LGD7i$-aM}Sni{?lsz zR$WXr7d72Lf5u+^rTy4;-&@~C?_I7UG=hgD*z}YH4HO?1xCN-9g@^EO1MmLWb@IFX z*85Ka=y2DK?gt(HTj@s~UDpP;CZ7Hgd*Ly{pjNjh-%oxMxVN(1``&~0}d7Pe`A-^SQZya!0^=*GF7 zSog-cCrV5m!ZmxsN+QDgVe=%1d#c|`UZkm){Nl4Np%KXx#na<=2-AVXdIUKti0{<% zj(VZ9wFJGx20%r;YkV9CWEh zWh*LGYCjzHBSA?OqEI&aza*Hr60TJjy;Bg zwjb>k+;i@|XU;w6-1G9EOeO<@!5)8a=3gd+{(%?DP?wSQ%Ya-*4DusAVhBbRA^e04 zRelm68Bs?mKSdx$?boQ#5Mnfoh|%6xmA>~^&8u1wtKE(GI=_ygFb%6`b@z$#x#2uw zRDRvtx(}TtI7iH(E97}JN4RvU(}sa_^w`BP8xm;j4&Xc!uxwxd@5!MuXGr4SI!5EyGulPuH_Yo?dZ}jMT%aVhB6F&< zgef8lG#8^eEP$|G#!%Td@>TeSNcdlW4JLR8cmalF)C|RF7%ih?^taSgDp(u7x30RS z=70j3U zdg$+|#S<>Iq-EHtU@{^~WR#1+6ucs-1(8im|K-nktYR8UHnc>^+DV2`JT}EKY%Ih| zCN7=~h%-DZ%)}$iX9$EKLlHs1>e8>K+y9E>Q*9`s8ogm1o7<)0^#CLj$l&%&dIl!FW4;O4t;V~CiW3XcjuKwhMo1Y4 z&h0Px6Zm+D!pzE};w>*Bey<}@*++s!#9ky1h7gtLMmHh)J|>lDe3#H8QLE6URiaM8 z!*ZrjRwby!I3G!?!jM0<4kcC?@=`lOJe3ANwF1$C3h2#SsZk}i!ZrW@jG+;Tywu{!v*a$TF{Ah z<-4?2`HE3rawC)>cuQJa(l+c+s5W4R_!ap>=L>qVMWNnwi(a`!fo}Bkufd*%v_8B` zq2GXM16K9NT6&=hxbEKIreHsjt5Mur;iM0PUu+#ejlW`<1yfoV?%c{7p>$QaYx}O^ zsp5X+ENxUcrA^|2$}Yj^Dqe3+6XAo($p%cDu+?J*u}9fUn-pF+TPrw$KKz^d$3zP>mdxXFgqiljHu$0`MXmteFD)$deq z>KRO_z2$abdp7_L0T)Xo#Zu)>)sJ$5+#fJ(f{o!C#j!#kojTP`k4M;`zyh@(1|t!= zTtcI4Oay`im#CaTKN6aup(OIb$h^pfpuhvK(B-bgyf_n&J;RBX$6dv<|H9-_NorgW zB@#-ZPoiSMC@VnCRA0Opzry3RCsCnDP!OPQVi3RtrAoXOKlu!2pOm%S4pr&{Jr(34 zEYnB3oKn5pGw$)ZJ-z|2XClyluHQS-|K^BCGM-7sL@vsDcs|Zc#!>IYgx5C&WoLXs zGH_r7&=Z$vO4dixOtLWn2uX&ELXA!&LDXavil-(hi1P_nQpFMymE>aLQ638=>8JR3 zG|2P8d5Mg$F^K|wSV5xrU~HO|)K})O%=0)8g5oie$C(;fC9u^cjTnz`f+$f_k$6y) zv@<~=D2hB!%jAN{O9sq5P!fPDU7}CMgxMs^rdUbuxe{U%B2+?G4Rj-HWQxbNMXsW< zN%>m&IxadL0I9VYCb)>Cg+2t{P2h33@`o@D!AT^32rOwyNMk%Hg2M#L%?$olESZL2 z1^iz0*1_Yz<8ft4(c7Dg?-%lJpb{M59aqsm&26i-j+K@pkBkrADIzpTjH>|9r&Z7^ zSk7SF1haad-a~w2w%%y_bm!gXJIznc-D{S{>&LDg`{?*`%M(j$p}A$bWobC$S+hBB zyOvxxJM*^AtgSO=>v}+DZHMyK<5}zRM{oXN_;*?yR!DK&oBOF z=ChdxbGhylIs3`HeK2bu%-M&(=*im0Glqi2e*KMWZ{#ek%d<}`?S+=kd`nNZrRPB~ z*K#y-s^D~GjB9oDa>t6LYsJ|0?KkGe7l?odm0!N^(V>U4zwdq6zTzCn)w)-V?(d2u zK7A!Lf}h+R?)BE9$F+Lz5%O_UgZCi$_#g$i&s8Pq0&=1V1SCTs0PS}&!UAp#1ZI=L z2u>le(gp%dJOoIUoLQnQq$h^Qof0Gm&*a%NzEVqES^4@$Jrj~{@T^a6OyRdf|D<{$(I1HGUXV z;mX6~21Ax*i;$A;^^$Sk;3t6u|LhP(VFEJ~#B&u%sEbA|Vf#v5{gt}@YqhmtZ&^{> z{$unJgok*6WbiHciO%u@sR&!az6*}OwA5#2Kk9vN min_area: + x, y, w, h = cv2.boundingRect(contour) + center_x = x + w // 2 + center_y = y + h // 2 + + tracks.append({ + 'bbox': [x, y, x + w, y + h], + 'center': [center_x, center_y], + 'area': area, + 'confidence': min(area / (min_area * 10), 1.0) + }) + + return tracks + + def _optical_flow_tracking(self, gray, previous_tracks): + """Lucas-Kanade optical flow tracking""" + tracks = [] + + if self.prev_frame is not None and previous_tracks: + # Extract previous points + prev_points = [] + for track in previous_tracks: + if 'center' in track: + prev_points.append([track['center']]) + + if prev_points: + prev_points = np.array(prev_points, dtype=np.float32) + + # Calculate optical flow + new_points, status, error = cv2.calcOpticalFlowPyrLK( + self.prev_frame, gray, prev_points, None + ) + + # Keep only good points + good_new = new_points[status == 1] + + for i, point in enumerate(good_new): + x, y = point.ravel() + tracks.append({ + 'center': [int(x), int(y)], + 'confidence': 0.8, + 'bbox': [int(x-20), int(y-20), int(x+20), int(y+20)] # Approximate bbox + }) + + return tracks + + def _contour_tracking(self, gray, min_area, threshold): + """Simple contour-based tracking""" + tracks = [] + + # Apply threshold + _, thresh = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) + + # Find contours + contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + for contour in contours: + area = cv2.contourArea(contour) + if area > min_area: + x, y, w, h = cv2.boundingRect(contour) + center_x = x + w // 2 + center_y = y + h // 2 + + tracks.append({ + 'bbox': [x, y, x + w, y + h], + 'center': [center_x, center_y], + 'area': area, + 'confidence': min(area / (min_area * 5), 1.0) + }) + + return tracks diff --git a/nodes/cv_model_loader.py b/nodes/cv_model_loader.py new file mode 100644 index 0000000..887a6eb --- /dev/null +++ b/nodes/cv_model_loader.py @@ -0,0 +1,57 @@ +""" +CV Model Loader Node for ComfyUI +Loads YOLO models for object detection +""" + +import os + +# Handle missing dependencies gracefully +try: + import torch + from ultralytics import YOLO + DEPENDENCIES_AVAILABLE = True +except ImportError as e: + DEPENDENCIES_AVAILABLE = False + MISSING_DEPS = str(e) + + +class CV_ModelLoader: + """Loads YOLO models for computer vision tasks""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_name": (["yolov8n.pt", "yolov8s.pt", "yolov8m.pt", "yolov8l.pt", "yolov8x.pt"], { + "default": "yolov8n.pt" + }), + }, + "optional": { + "custom_model_path": ("STRING", { + "default": "", + "multiline": False + }), + } + } + + RETURN_TYPES = ("CV_MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "load_model" + CATEGORY = "CV/Models" + + def load_model(self, model_name, custom_model_path=""): + """Load YOLO model""" + if not DEPENDENCIES_AVAILABLE: + raise RuntimeError(f"Missing dependencies: {MISSING_DEPS}. Please install requirements: pip install torch ultralytics") + + try: + # Use custom path if provided, otherwise use default model + model_path = custom_model_path if custom_model_path else model_name + + # Load YOLO model + model = YOLO(model_path) + + return (model,) + + except Exception as e: + raise RuntimeError(f"Failed to load model: {str(e)}") diff --git a/nodes/cv_object_detector.py b/nodes/cv_object_detector.py new file mode 100644 index 0000000..8495b9c --- /dev/null +++ b/nodes/cv_object_detector.py @@ -0,0 +1,95 @@ +""" +CV Object Detector Node for ComfyUI +Performs object detection using YOLO models +""" + +# Handle missing dependencies gracefully +try: + import torch + import numpy as np + from PIL import Image + DEPENDENCIES_AVAILABLE = True +except ImportError as e: + DEPENDENCIES_AVAILABLE = False + MISSING_DEPS = str(e) + + +class CV_ObjectDetector: + """Detects objects in images using YOLO models""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("CV_MODEL",), + "image": ("IMAGE",), + "confidence": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "iou_threshold": ("FLOAT", { + "default": 0.45, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + } + } + + RETURN_TYPES = ("IMAGE", "CV_DETECTIONS") + RETURN_NAMES = ("image", "detections") + FUNCTION = "detect_objects" + CATEGORY = "CV/Detection" + + def detect_objects(self, model, image, confidence, iou_threshold): + """Perform object detection on input image""" + if not DEPENDENCIES_AVAILABLE: + raise RuntimeError(f"Missing dependencies: {MISSING_DEPS}. Please install requirements: pip install torch ultralytics opencv-python") + + try: + # Convert ComfyUI image tensor to PIL Image + if isinstance(image, torch.Tensor): + # ComfyUI images are in format (batch, height, width, channels) + if image.dim() == 4: + image = image[0] # Take first image from batch + # Convert from tensor to numpy + img_np = (image.cpu().numpy() * 255).astype(np.uint8) + img_pil = Image.fromarray(img_np) + else: + img_pil = image + + # Run inference + results = model(img_pil, conf=confidence, iou=iou_threshold) + + # Extract detection data + detections = [] + if len(results) > 0 and results[0].boxes is not None: + boxes = results[0].boxes + for i in range(len(boxes)): + box = boxes.xyxy[i].cpu().numpy() # [x1, y1, x2, y2] + conf = boxes.conf[i].cpu().numpy() + cls = int(boxes.cls[i].cpu().numpy()) + + detections.append({ + 'bbox': box.tolist(), + 'confidence': float(conf), + 'class': cls, + 'class_name': model.names[cls] if hasattr(model, 'names') else str(cls) + }) + + # Convert back to ComfyUI tensor format + if isinstance(image, torch.Tensor): + output_image = image + else: + # Convert PIL to tensor if needed + img_tensor = torch.from_numpy(np.array(img_pil)).float() / 255.0 + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + output_image = img_tensor + + return (output_image, detections) + + except Exception as e: + raise RuntimeError(f"Object detection failed: {str(e)}") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..40e0f62 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +opencv-python +ultralytics +torch +torchvision +numpy +Pillow \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..dd7ee44 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1 @@ +# Utils package diff --git a/utils/opencv_helpers.py b/utils/opencv_helpers.py new file mode 100644 index 0000000..e49ae6a --- /dev/null +++ b/utils/opencv_helpers.py @@ -0,0 +1,103 @@ +""" +Utility functions for OpenCV operations +""" + +import cv2 +import numpy as np +import torch +from PIL import Image + + +def tensor_to_opencv(tensor): + """Convert ComfyUI tensor to OpenCV format""" + if isinstance(tensor, torch.Tensor): + if tensor.dim() == 4: + tensor = tensor[0] # Take first image from batch + img_np = (tensor.cpu().numpy() * 255).astype(np.uint8) + # Convert RGB to BGR for OpenCV + if len(img_np.shape) == 3: + img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + return img_np + return tensor + + +def opencv_to_tensor(img_np): + """Convert OpenCV format to ComfyUI tensor""" + # Convert BGR to RGB + if len(img_np.shape) == 3: + img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB) + + # Normalize and convert to tensor + img_tensor = torch.from_numpy(img_np).float() / 255.0 + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + + return img_tensor + + +def tensor_to_pil(tensor): + """Convert ComfyUI tensor to PIL Image""" + if isinstance(tensor, torch.Tensor): + if tensor.dim() == 4: + tensor = tensor[0] # Take first image from batch + img_np = (tensor.cpu().numpy() * 255).astype(np.uint8) + return Image.fromarray(img_np) + return tensor + + +def pil_to_tensor(img_pil): + """Convert PIL Image to ComfyUI tensor""" + img_np = np.array(img_pil).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_np) + if img_tensor.dim() == 3: + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + return img_tensor + + +def normalize_bbox(bbox, img_width, img_height): + """Normalize bounding box coordinates to 0-1 range""" + x1, y1, x2, y2 = bbox + return [ + x1 / img_width, + y1 / img_height, + x2 / img_width, + y2 / img_height + ] + + +def denormalize_bbox(bbox, img_width, img_height): + """Denormalize bounding box coordinates from 0-1 range to pixel coordinates""" + x1, y1, x2, y2 = bbox + return [ + int(x1 * img_width), + int(y1 * img_height), + int(x2 * img_width), + int(y2 * img_height) + ] + + +def resize_image_keep_aspect(image, target_size): + """Resize image while keeping aspect ratio""" + h, w = image.shape[:2] + target_w, target_h = target_size + + # Calculate scale factor + scale = min(target_w / w, target_h / h) + + # Calculate new dimensions + new_w = int(w * scale) + new_h = int(h * scale) + + # Resize image + resized = cv2.resize(image, (new_w, new_h)) + + # Create padded image + result = np.zeros((target_h, target_w, 3), dtype=image.dtype) + + # Center the resized image + y_offset = (target_h - new_h) // 2 + x_offset = (target_w - new_w) // 2 + + result[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized + + return result, scale, (x_offset, y_offset) diff --git a/utils/yolo_utils.py b/utils/yolo_utils.py new file mode 100644 index 0000000..6f28e20 --- /dev/null +++ b/utils/yolo_utils.py @@ -0,0 +1,132 @@ +""" +Utility functions for YOLO operations +""" + +import torch +import numpy as np +from ultralytics import YOLO + + +def load_yolo_model(model_path): + """Load YOLO model safely""" + try: + model = YOLO(model_path) + return model + except Exception as e: + raise RuntimeError(f"Failed to load YOLO model from {model_path}: {str(e)}") + + +def filter_detections_by_class(detections, allowed_classes): + """Filter detections by class names or IDs""" + if not allowed_classes: + return detections + + filtered = [] + for detection in detections: + class_id = detection.get('class', -1) + class_name = detection.get('class_name', '') + + if class_id in allowed_classes or class_name in allowed_classes: + filtered.append(detection) + + return filtered + + +def filter_detections_by_confidence(detections, min_confidence): + """Filter detections by confidence threshold""" + return [d for d in detections if d.get('confidence', 0) >= min_confidence] + + +def filter_detections_by_area(detections, min_area=None, max_area=None): + """Filter detections by bounding box area""" + filtered = [] + + for detection in detections: + bbox = detection['bbox'] + x1, y1, x2, y2 = bbox + area = (x2 - x1) * (y2 - y1) + + if min_area is not None and area < min_area: + continue + if max_area is not None and area > max_area: + continue + + filtered.append(detection) + + return filtered + + +def nms_detections(detections, iou_threshold=0.45): + """Apply Non-Maximum Suppression to detections""" + if not detections: + return detections + + # Convert to format for NMS + boxes = [] + scores = [] + + for detection in detections: + bbox = detection['bbox'] + boxes.append(bbox) + scores.append(detection.get('confidence', 1.0)) + + boxes = torch.tensor(boxes, dtype=torch.float32) + scores = torch.tensor(scores, dtype=torch.float32) + + # Apply NMS + keep_indices = torch.ops.torchvision.nms(boxes, scores, iou_threshold) + + # Return filtered detections + return [detections[i] for i in keep_indices] + + +def get_yolo_class_names(): + """Get standard YOLO class names""" + return [ + 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', + 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', + 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', + 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', + 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', + 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', + 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', + 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', + 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', + 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', + 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', + 'toothbrush' + ] + + +def convert_yolo_to_detection_format(results, img_width, img_height): + """Convert YOLO results to standardized detection format""" + detections = [] + + if len(results) > 0 and results[0].boxes is not None: + boxes = results[0].boxes + + for i in range(len(boxes)): + # Get box coordinates + box = boxes.xyxy[i].cpu().numpy() # [x1, y1, x2, y2] + conf = float(boxes.conf[i].cpu().numpy()) + cls = int(boxes.cls[i].cpu().numpy()) + + # Get class name if available + class_names = get_yolo_class_names() + class_name = class_names[cls] if cls < len(class_names) else f"class_{cls}" + + # Create detection dictionary + detection = { + 'bbox': box.tolist(), + 'confidence': conf, + 'class': cls, + 'class_name': class_name, + 'center': [(box[0] + box[2]) / 2, (box[1] + box[3]) / 2], + 'width': box[2] - box[0], + 'height': box[3] - box[1], + 'area': (box[2] - box[0]) * (box[3] - box[1]) + } + + detections.append(detection) + + return detections