# Implementation Plan: Compass Image Loader

This document provides a step-by-step implementation plan for the `CompassImageLoader` ComfyUI custom node. Each task is a concrete, actionable item with a checkbox.

---

## Phase 1: Project Bootstrap & File Structure

**Goal:** Set up the custom node package so ComfyUI recognizes and registers the node.

- [ ] **Create `__init__.py`**
  - [ ] Import `CompassImageLoader` from `compass_image_loader`.
  - [ ] Define `NODE_CLASS_MAPPINGS = {"CompassImageLoader": CompassImageLoader}`.
  - [ ] Define `NODE_DISPLAY_NAME_MAPPINGS = {"CompassImageLoader": "Compass Image Loader"}`.

- [ ] **Create `compass_image_loader.py`**
  - [ ] Define the `CompassImageLoader` class.
  - [ ] Set `CATEGORY = "image/loaders"`.

---

## Phase 2: Backend Node Implementation

**Goal:** Implement the Python class that ComfyUI executes.

### 2.1 Input Types & Discovery

- [ ] Implement `CompassImageLoader.INPUT_TYPES(cls)` classmethod.
  - [ ] Implement `_discover_directories()` helper that:
    - [ ] Gets the base path via `folder_paths.get_input_directory()`.
    - [ ] Recursively walks the base path using `os.walk`.
    - [ ] Identifies valid direction names: `{"n", "ne", "e", "se", "s", "sw", "w", "nw"}`.
    - [ ] Identifies valid modality names: `{"image", "depth", "openpose"}`.
    - [ ] For each directory, if it contains **any** valid direction subfolder **or** any valid modality subfolder, add its relative path (from `input/`) to the candidate list.
    - [ ] Return sorted, deduplicated list of strings.
  - [ ] Return dictionary with:
    - `required`:
      - `directory`: `(directories_list,)`
      - `direction`: `(["", "n", "ne", "e", "se", "s", "sw", "w", "nw"],)`
      - `modality`: `(["image", "depth", "openpose"],)`
      - `frame`: `("STRING", {"default": ""})`
      - `width`: `("INT", {"default": 0, "min": 0, "max": 16384, "step": 1})`
      - `height`: `("INT", {"default": 0, "min": 0, "max": 16384, "step": 1})`

### 2.2 Return Types & Function Signature

- [ ] Set `RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "INT")`.
- [ ] Set `RETURN_NAMES = ("IMAGE", "path", "width", "height", "frame_count")`.
- [ ] Set `FUNCTION = "load_images"`.
- [ ] Set `OUTPUT_NODE = False`.

### 2.3 Core Loading Logic

- [ ] Implement `load_images(self, directory, direction, modality, frame, width, height)`.

#### Path Resolution
- [ ] Resolve `base_dir = folder_paths.get_input_directory()`.
- [ ] Construct `target_dir`:
  - [ ] If `direction.strip()`: `os.path.join(base_dir, directory, direction, modality)`.
  - [ ] Else: `os.path.join(base_dir, directory, modality)`.
- [ ] Verify `target_dir` exists; if not, raise `RuntimeError`.

#### Image Discovery
- [ ] List files in `target_dir` with `os.listdir`.
- [ ] Filter to supported image extensions: `png`, `jpg`, `jpeg`, `webp`, `bmp`, `gif`, `tiff`.
- [ ] Sort alphabetically: `sorted(files)`.
- [ ] If list is empty, raise `RuntimeError`.

#### Frame Selection
- [ ] If `frame` is `None` or `frame.strip() == ""`:
  - [ ] Select all images (`selected_files = all_files`).
  - [ ] Set `output_path = target_dir`.
- [ ] Else:
  - [ ] Try `index = int(frame.strip())`.
  - [ ] If `index < 0` or `index >= len(all_files)`, raise `RuntimeError`.
  - [ ] Select single file: `selected_files = [all_files[index]]`.
  - [ ] Set `output_path = os.path.join(target_dir, all_files[index])`.

#### Image Processing Loop
- [ ] Initialize empty list `tensors = []`.
- [ ] For each file in `selected_files`:
  - [ ] Open with `Image.open(path)`.
  - [ ] Convert to RGB: `image = image.convert("RGB")`.
  - [ ] Read original dimensions: `orig_w, orig_h = image.size`.
  - [ ] Determine final dimensions and resize:
    - [ ] If `width == 0 and height == 0`:
      - [ ] `final_w, final_h = orig_w, orig_h` (no resize).
    - [ ] Elif `width > 0 and height == 0`:
      - [ ] `final_w = width`.
      - [ ] `final_h = int(orig_h * (width / orig_w))`.
      - [ ] `image = image.resize((final_w, final_h), Image.Resampling.LANCZOS)`.
    - [ ] Elif `height > 0 and width == 0`:
      - [ ] `final_h = height`.
      - [ ] `final_w = int(orig_w * (height / orig_h))`.
      - [ ] `image = image.resize((final_w, final_h), Image.Resampling.LANCZOS)`.
    - [ ] Else (both non-zero):
      - [ ] `scale = max(width / orig_w, height / orig_h)`.
      - [ ] `new_w = int(orig_w * scale)`.
      - [ ] `new_h = int(orig_h * scale)`.
      - [ ] `image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)`.
      - [ ] Compute crop box:
        - [ ] `left = (new_w - width) // 2`
        - [ ] `top = (new_h - height) // 2`
        - [ ] `right = left + width`
        - [ ] `bottom = top + height`
      - [ ] `image = image.crop((left, top, right, bottom))`.
      - [ ] `final_w, final_h = width, height`.
  - [ ] Convert to NumPy array:
    - [ ] `np_image = np.array(image).astype(np.float32) / 255.0`
  - [ ] Convert to Torch tensor:
    - [ ] `tensor = torch.from_numpy(np_image)[None,]`  # shape [1, H, W, 3]
  - [ ] Append to `tensors`.

#### Batch Assembly
- [ ] If `len(tensors) == 1`:
  - [ ] `image_batch = tensors[0]`
- [ ] Else:
  - [ ] `image_batch = torch.cat(tensors, dim=0)`

#### Return
- [ ] Return `(image_batch, output_path, final_w, final_h, len(selected_files))`.

### 2.4 Change Detection

- [ ] Implement `IS_CHANGED(self, directory, direction, modality, frame, width, height)`:
  - [ ] Compute `target_dir` as above.
  - [ ] If frame is empty, hash the sorted filenames list + mtime.
  - [ ] If frame is specified, hash the single file content.
  - [ ] Include `width` and `height` in the hash (so resizing triggers re-execution).

---

## Phase 3: Frontend Extension

**Goal:** Provide an enhanced UX for the `directory` combo widget (typing to filter / autocomplete).

- [ ] **Create `js/compass_image_loader.js`**
  - [ ] Import `app` from `scripts/app.js`.
  - [ ] Import `api` from `scripts/api.js` (for event listeners if needed).

- [ ] **Register Extension**
  - [ ] Call `app.registerExtension({ name: "CompassImageLoader", ... })`.

- [ ] **Override `beforeRegisterNodeDef`**
  - [ ] Check `passedNodeData.name === "CompassImageLoader"`.
  - [ ] Save reference to original `onNodeCreated`.
  - [ ] Replace with a wrapper that:
    - [ ] Calls the original `onNodeCreated`.
    - [ ] Finds the `directory` widget by name.
    - [ ] Stores the original full list of combo values on the widget instance (e.g., `widget._all_values`).
    - [ ] Attaches a keyup/input listener to the widget's DOM element (`widget.inputEl` or equivalent) to implement filtering.
    - [ ] When the user types, filter `_all_values` case-insensitively and update `widget.options.values`.
    - [ ] If the current `widget.value` is not in the filtered list, leave it as-is (allow free-form or preserve selection).

- [ ] **Handle Refresh**
  - [ ] Listen for `fresh-node-defs` event from `api`.
  - [ ] On event, iterate over all nodes in `app.graph._nodes`.
  - [ ] For each `CompassImageLoader` node:
    - [ ] Find the `directory` widget.
    - [ ] Update `widget._all_values` and `widget.options.values` with the new list from the refreshed node definition.
    - [ ] If the currently selected value still exists in the new list, keep it.
    - [ ] If it does not exist, clear it to `""` (or first available).

- [ ] **State Preservation**
  - [ ] Ensure the directory value is correctly serialized in the workflow JSON (ComfyUI handles this automatically for widgets, but verify the filtered list does not override the saved value on `configure`).

---

## Phase 4: Integration & Packaging

- [ ] Verify file layout:
  ```
  ComfyUI/custom_nodes/ComfyUI-compass-paths/
    __init__.py
    compass_image_loader.py
    js/
      compass_image_loader.js
  ```
- [ ] Ensure `__init__.py` is minimal and correct.
- [ ] Restart ComfyUI and verify the node appears in the node search under `image/loaders`.

---

## Phase 5: Testing & Validation

Run the following manual tests in the ComfyUI interface.

- [ ] **Test 1: Single Frame Load**
  - [ ] Set `directory`, `direction=s`, `modality=depth`, `frame=0`, `width=0`, `height=0`.
  - [ ] Queue. Verify output `IMAGE` shape is `[1, H, W, 3]`, `frame_count=1`, `path` is the full file path.

- [ ] **Test 2: Batch Load**
  - [ ] Leave `frame` blank.
  - [ ] Queue. Verify `IMAGE` shape is `[N, H, W, 3]` where `N > 1`, `frame_count=N`, `path` is the directory path.

- [ ] **Test 3: Cover + Crop Resize**
  - [ ] Set `width=512`, `height=768` on a non-square source image.
  - [ ] Verify output `width=512`, `height=768`, and image appears center-cropped.

- [ ] **Test 4: Proportional Resize (Width Only)**
  - [ ] Set `width=512`, `height=0`.
  - [ ] Verify output width=512 and height is proportionally scaled (no cropping).

- [ ] **Test 5: Proportional Resize (Height Only)**
  - [ ] Set `width=0`, `height=768`.
  - [ ] Verify output height=768 and width is proportionally scaled (no cropping).

- [ ] **Test 6: No Direction**
  - [ ] Set `direction` to blank.
  - [ ] Verify node loads from `input/<directory>/<modality>/` (no direction folder).

- [ ] **Test 7: Error — Missing Directory**
  - [ ] Set `directory` to a non-existent path.
  - [ ] Queue. Verify a clear `RuntimeError` is raised.

- [ ] **Test 8: Error — Empty Image Folder**
  - [ ] Point to a valid directory/modality path that has no images.
  - [ ] Queue. Verify a clear `RuntimeError` is raised.

- [ ] **Test 9: Error — Out of Bounds Frame**
  - [ ] Set `frame` to a number larger than the number of images minus one.
  - [ ] Queue. Verify a clear `RuntimeError` is raised.

- [ ] **Test 10: Refresh Updates Combo**
  - [ ] Add a new valid folder under `input/` while ComfyUI is running.
  - [ ] Click the ComfyUI Refresh button.
  - [ ] Open the node's `directory` dropdown. Verify the new folder appears.

- [ ] **Test 11: Autocomplete Filtering**
  - [ ] Click the `directory` combo and type a substring (e.g., `walk`).
  - [ ] Verify the dropdown filters to show only matching entries.

---

## Future Enhancements (Out of Scope for Initial Implementation)

- Support for recursive image loading across subdirectories.
- Caching of directory scans to improve performance on very large `input/` trees.
- Support for additional modalities beyond `image`, `depth`, `openpose` via a configuration file.
- Drag-and-drop folder path into the directory widget.