diff --git a/image copy.png b/image copy.png new file mode 100644 index 0000000..c24c539 Binary files /dev/null and b/image copy.png differ diff --git a/src/atm/calibrate.py b/src/atm/calibrate.py index 23864c3..48a4df9 100644 --- a/src/atm/calibrate.py +++ b/src/atm/calibrate.py @@ -66,31 +66,37 @@ def write_config(data: dict, out_dir: Path) -> Path: # Window capture (Windows live) / screenshot fallback # --------------------------------------------------------------------------- -def _capture_window(title_substr: str): - """Screenshot the first window whose title CONTAINS *title_substr* (case-insensitive). - - Windows-only (mss+pygetwindow). Returns a PIL RGB Image. - """ - import mss # type: ignore[import-untyped] +def _find_candidate_windows(title_substr: str): + """Return list of (title, width, height, window_obj) matching substring, sorted by area desc.""" import pygetwindow as gw # type: ignore[import-untyped] - from PIL import Image needle = title_substr.lower() - all_titles = [t for t in gw.getAllTitles() if t] - matching = [t for t in all_titles if needle in t.lower()] - if not matching: - preview = "\n ".join(sorted(set(all_titles))[:40]) - raise RuntimeError( - f"No window title contains {title_substr!r} (case-insensitive).\n" - f"Visible windows (first 40):\n {preview}" - ) - win = gw.getWindowsWithTitle(matching[0])[0] - try: - win.activate() - except Exception: - pass # activate may fail on some windows; continue — mss still captures by region - time.sleep(0.3) + out = [] + seen = set() + for t in gw.getAllTitles(): + if not t or needle not in t.lower(): + continue + if t in seen: + continue + seen.add(t) + wins = gw.getWindowsWithTitle(t) + if not wins: + continue + w = wins[0] + try: + width = int(w.width); height = int(w.height) + except Exception: + continue + if width < 10 or height < 10: + continue + out.append((t, width, height, w)) + out.sort(key=lambda r: r[1] * r[2], reverse=True) + return out + +def _capture_via_mss(win): + import mss # type: ignore[import-untyped] + from PIL import Image with mss.mss() as sct: mon = { "top": int(win.top), "left": int(win.left), @@ -100,6 +106,157 @@ def _capture_window(title_substr: str): return Image.frombytes("RGB", shot.size, shot.rgb) +def _capture_via_printwindow(win): + """Use Win32 PrintWindow with PW_RENDERFULLCONTENT. Reliable for GPU-accelerated apps.""" + import win32gui # type: ignore[import-untyped] + import win32ui # type: ignore[import-untyped] + from ctypes import windll + from PIL import Image + + hwnd = win._hWnd # pygetwindow exposes this + rect = win32gui.GetWindowRect(hwnd) + w = rect[2] - rect[0] + h = rect[3] - rect[1] + if w <= 0 or h <= 0: + raise RuntimeError("window has zero size") + + hwnd_dc = win32gui.GetWindowDC(hwnd) + mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc) + save_dc = mfc_dc.CreateCompatibleDC() + save_bitmap = win32ui.CreateBitmap() + save_bitmap.CreateCompatibleBitmap(mfc_dc, w, h) + save_dc.SelectObject(save_bitmap) + + PW_RENDERFULLCONTENT = 2 + result = windll.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), PW_RENDERFULLCONTENT) + bmp_info = save_bitmap.GetInfo() + bmp_bits = save_bitmap.GetBitmapBits(True) + img = Image.frombuffer( + "RGB", (bmp_info["bmWidth"], bmp_info["bmHeight"]), + bmp_bits, "raw", "BGRX", 0, 1, + ) + + win32gui.DeleteObject(save_bitmap.GetHandle()) + save_dc.DeleteDC() + mfc_dc.DeleteDC() + win32gui.ReleaseDC(hwnd, hwnd_dc) + + if result != 1: + raise RuntimeError(f"PrintWindow returned {result}") + return img + + +def _image_looks_blank(pil_img, dark_thresh: int = 10) -> bool: + """Heuristic: image is effectively blank if near-black or near-uniform.""" + import numpy as np + arr = np.array(pil_img.resize((64, 36))).astype(np.int16) + if arr.mean() < dark_thresh: + return True + return arr.std() < 3.0 + + +def _capture_window(title_substr: str): + """Capture a window by title substring. Tries PrintWindow first (reliable for GPU + apps), falls back to mss. If multiple windows match, picks the largest by area. + + Returns a PIL RGB Image. + """ + import pygetwindow as gw # type: ignore[import-untyped] + + candidates = _find_candidate_windows(title_substr) + if not candidates: + all_titles = sorted({t for t in gw.getAllTitles() if t}) + preview = "\n ".join(all_titles[:40]) + raise RuntimeError( + f"No window title contains {title_substr!r} (case-insensitive).\n" + f"Visible windows (first 40):\n {preview}" + ) + + # If multiple candidates, ask user to pick (Tk listbox) + if len(candidates) > 1: + picked = _pick_window_dialog(candidates) + else: + picked = candidates[0] + title, w_, h_, win = picked + + try: + win.activate() + except Exception: + pass + time.sleep(0.3) + + errors = [] + img = None + for fn, name in ((_capture_via_printwindow, "PrintWindow"), + (_capture_via_mss, "mss")): + try: + candidate_img = fn(win) + if _image_looks_blank(candidate_img): + errors.append(f"{name}: returned blank/uniform image") + continue + img = candidate_img + break + except Exception as e: + errors.append(f"{name}: {e}") + + if img is None: + raise RuntimeError( + f"Capture of window {title!r} ({w_}x{h_}) failed:\n " + "\n ".join(errors) + ) + + # Save debug copy so user can verify + try: + debug_path = Path("logs") / f"calibrate_capture_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" + debug_path.parent.mkdir(exist_ok=True) + img.save(debug_path) + except Exception: + pass + + return img + + +def _pick_window_dialog(candidates): + """Show a Tk listbox for the user to pick a window.""" + import tkinter as tk + + root = tk.Tk() + root.title("Pick window") + root.geometry("600x400+100+100") + + tk.Label(root, text="Multiple windows match. Pick the main TradeStation chart:", + font=("Arial", 11), pady=6).pack(fill="x") + + frame = tk.Frame(root) + frame.pack(fill="both", expand=True, padx=8, pady=4) + scroll = tk.Scrollbar(frame) + scroll.pack(side="right", fill="y") + + lb = tk.Listbox(frame, font=("Consolas", 10), yscrollcommand=scroll.set) + lb.pack(side="left", fill="both", expand=True) + scroll.config(command=lb.yview) + + for title, w, h, _win in candidates: + lb.insert(tk.END, f"[{w}x{h}] {title}") + lb.selection_set(0) + + choice = {"idx": None} + + def _ok(): + sel = lb.curselection() + if sel: + choice["idx"] = sel[0] + root.destroy() + + tk.Button(root, text="OK", command=_ok, bg="#2ecc71", fg="white", + font=("Arial", 11), pady=4).pack(pady=6) + lb.bind("", lambda e: _ok()) + + root.mainloop() + if choice["idx"] is None: + raise ValueError("Window selection cancelled") + return candidates[choice["idx"]] + + # --------------------------------------------------------------------------- # Wizard # ---------------------------------------------------------------------------