master_adapt_detect/example_code/composite_image_finder.py
2025-10-01 14:32:55 -05:00

311 lines
16 KiB
Python

import sys
import subprocess
import tkinter as tk
from tkinter import messagebox, filedialog, scrolledtext, ttk
import json
import csv
import threading
import time
import os
try:
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
except ImportError:
root = tk.Tk()
root.withdraw()
if messagebox.askyesno("Dependency Error", "Required libraries are not installed. Would you like to try and install them now?"):
try:
subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python", "numpy", "Pillow"])
messagebox.showinfo("Success", "Dependencies installed successfully. Please restart the application.")
except Exception as e:
messagebox.showerror("Installation Failed", f"Could not install dependencies. Please run 'pip install -r requirements.txt' manually.\n\nError: {e}")
root.destroy()
sys.exit()
class ToolTip:
def __init__(self, widget, text):
self.widget = widget
self.text = text
self.tooltip = None
self.widget.bind("<Enter>", self.enter)
self.widget.bind("<Leave>", self.leave)
def enter(self, event=None):
x, y, _, _ = self.widget.bbox("insert")
x += self.widget.winfo_rootx() + 25
y += self.widget.winfo_rooty() + 25
self.tooltip = tk.Toplevel(self.widget)
self.tooltip.wm_overrideredirect(True)
self.tooltip.wm_geometry(f"+{x}+{y}")
label = tk.Label(self.tooltip, text=self.text, background="#ffffe0", relief="solid", borderwidth=1, wraplength=200)
label.pack()
def leave(self, event=None):
if self.tooltip:
self.tooltip.destroy()
self.tooltip = None
class MasterImageFinderApp:
def __init__(self, root):
self.root = root
self.root.title("Master Image Finder")
self.root.geometry("800x700")
self.layouts_path = tk.StringVar()
self.masters_path = tk.StringVar()
self.upscale = tk.BooleanVar()
self.denoise = tk.BooleanVar()
self.sharpen = tk.BooleanVar()
self.contrast = tk.BooleanVar()
tk.Label(root, text="Layouts Folder:").pack(pady=5)
tk.Entry(root, textvariable=self.layouts_path, width=100).pack(pady=5)
tk.Button(root, text="Select Layouts Folder", command=self.select_layouts_folder).pack(pady=5)
tk.Label(root, text="Master Images Folder:").pack(pady=5)
tk.Entry(root, textvariable=self.masters_path, width=100).pack(pady=5)
tk.Button(root, text="Select Master Images Folder", command=self.select_masters_folder).pack(pady=5)
enhancement_frame = tk.LabelFrame(root, text="Advanced Enhancement Options", padx=10, pady=10)
enhancement_frame.pack(pady=10, padx=10, fill="x")
upscale_check = tk.Checkbutton(enhancement_frame, text="Smart Upscaling", variable=self.upscale)
upscale_check.grid(row=0, column=0, sticky="w")
ToolTip(upscale_check, "Enlarges small images to improve feature detection. Best for images under 400x400px.")
denoise_check = tk.Checkbutton(enhancement_frame, text="Denoising", variable=self.denoise)
denoise_check.grid(row=0, column=1, sticky="w", padx=10)
ToolTip(denoise_check, "Removes digital noise and compression artifacts. Can be slow on large images.")
sharpen_check = tk.Checkbutton(enhancement_frame, text="Sharpening", variable=self.sharpen)
sharpen_check.grid(row=1, column=0, sticky="w")
ToolTip(sharpen_check, "Enhances edges and fine details. Very fast.")
contrast_check = tk.Checkbutton(enhancement_frame, text="Contrast Enhancement", variable=self.contrast)
contrast_check.grid(row=1, column=1, sticky="w", padx=10)
ToolTip(contrast_check, "Improves local contrast, making features in dark or washed-out areas more distinct.")
self.run_button = tk.Button(root, text="Find Matches", command=self.run_finder_thread)
self.run_button.pack(pady=20)
self.progress = ttk.Progressbar(root, orient="horizontal", length=780, mode="determinate")
self.progress.pack(pady=10)
self.log_area = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=100, height=15)
self.log_area.pack(pady=10, padx=10)
def select_layouts_folder(self):
self.layouts_path.set(filedialog.askdirectory())
def select_masters_folder(self):
self.masters_path.set(filedialog.askdirectory())
def log(self, message):
self.root.after(0, self._log, message)
def _log(self, message):
self.log_area.insert(tk.END, message + "\n")
self.log_area.see(tk.END)
def update_progress(self, value):
self.root.after(0, self.progress.config, {'value': value})
def run_finder_thread(self):
self.run_button.config(state=tk.DISABLED)
self.log_area.delete(1.0, tk.END)
self.progress['value'] = 0
thread = threading.Thread(target=self.run_finder)
thread.start()
def run_finder(self):
start_time = time.time()
layouts_dir = self.layouts_path.get()
masters_dir = self.masters_path.get()
if not layouts_dir or not masters_dir:
self.log("Error: Please select both folders.")
self.run_button.config(state=tk.NORMAL)
return
output_dir = os.path.join(layouts_dir, "reports")
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
self.log(f"Created output directory at {output_dir}")
try:
results = self.find_matches(layouts_dir, masters_dir)
self.create_html_report(results, output_dir, layouts_dir, masters_dir)
end_time = time.time()
total_matches = sum(1 for item in results if item['found'])
self.log("\n--- Process Complete! ---")
self.log(f"Found matches for {total_matches} out of {len(results)} layout images.")
self.log(f"Total time: {end_time - start_time:.2f} seconds")
self.log(f"Reports saved in: {output_dir}")
except Exception as e:
self.log(f"An error occurred: {e}")
finally:
self.run_button.config(state=tk.NORMAL)
def find_matches(self, layouts_path, masters_path, min_good_matches=10, inlier_threshold_ratio=0.5):
akaze = cv2.AKAZE_create()
bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
sharpen_kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
layout_images = [f for f in os.listdir(layouts_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
master_images = [f for f in os.listdir(masters_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
results = []
master_descriptors = {}
self.log("Preprocessing master images...")
for i, master_image_name in enumerate(master_images):
self.log(f" - Preprocessing {i+1}/{len(master_images)}: {master_image_name}")
master_image_path = os.path.join(masters_path, master_image_name)
master_img = cv2.imread(master_image_path, cv2.IMREAD_GRAYSCALE)
if master_img is None: continue
kp, des = akaze.detectAndCompute(master_img, None)
if des is not None: master_descriptors[master_image_name] = (kp, des)
total_layouts = len(layout_images)
self.log("\nProcessing layout images...")
for i, layout_image_name in enumerate(layout_images):
self.update_progress((i / total_layouts) * 100)
self.log(f" - Processing {i+1}/{total_layouts}: {layout_image_name}")
layout_image_path = os.path.join(layouts_path, layout_image_name)
layout_img_gray = cv2.imread(layout_image_path, cv2.IMREAD_GRAYSCALE)
if layout_img_gray is None:
self.log(f" - Could not read layout image.")
continue
enhancements_applied = []
if self.upscale.get() and (layout_img_gray.shape[0] < 400 or layout_img_gray.shape[1] < 400):
layout_img_gray = cv2.resize(layout_img_gray, (0,0), fx=2.0, fy=2.0, interpolation=cv2.INTER_LANCZOS4)
enhancements_applied.append("Upscaled")
if self.denoise.get():
layout_img_gray = cv2.fastNlMeansDenoising(layout_img_gray, None, 10, 7, 21)
enhancements_applied.append("Denoised")
if self.sharpen.get():
layout_img_gray = cv2.filter2D(layout_img_gray, -1, sharpen_kernel)
enhancements_applied.append("Sharpened")
if self.contrast.get():
layout_img_gray = clahe.apply(layout_img_gray)
enhancements_applied.append("Contrast Enhanced")
if enhancements_applied:
self.log(f" - Applied: {', '.join(enhancements_applied)}")
kp1, des1 = akaze.detectAndCompute(layout_img_gray, None)
if des1 is None:
self.log(f" - No features found.")
results.append({"layout": layout_image_name, "masters": [], "found": False, "enhancements": enhancements_applied})
continue
all_possible_matches = []
for master_image_name, (kp2, des2) in master_descriptors.items():
matches = bf.knnMatch(des1, des2, k=2)
good_matches = [m for m, n in matches if len(matches) > 1 and len(matches[0]) > 1 and m.distance < 0.75 * n.distance]
if len(good_matches) > min_good_matches:
src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
if mask is not None:
inliers = np.sum(mask)
if inliers > min_good_matches:
all_possible_matches.append({"master": master_image_name, "inliers": int(inliers)})
if not all_possible_matches:
results.append({"layout": layout_image_name, "masters": [], "found": False, "enhancements": enhancements_applied})
continue
best_match = max(all_possible_matches, key=lambda x: x['inliers'])
max_inliers = best_match['inliers']
confident_matches = [best_match]
for match in all_possible_matches:
if match != best_match and match['inliers'] > max_inliers * inlier_threshold_ratio:
confident_matches.append(match)
if confident_matches:
self.log(f" - Found {len(confident_matches)} confident master image(s).")
results.append({"layout": layout_image_name, "masters": confident_matches, "found": True, "enhancements": enhancements_applied})
else:
results.append({"layout": layout_image_name, "masters": [], "found": False, "enhancements": enhancements_applied})
self.update_progress(100)
return results
def create_html_report(self, data, output_path, layouts_abs_path, masters_abs_path):
report_path = os.path.join(output_path, 'report.html')
total_matches = sum(1 for item in data if item['found'])
total_layouts = len(data)
html = f"""
<!DOCTYPE html><html lang="en"><head><meta charset="UTF-8"><title>Image Match Report</title>
<style>
body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; margin: 20px; background-color: #f9f9f9; }}
.container {{ max-width: 1600px; margin: auto; }}
h1, h2 {{ text-align: center; color: #333; border-bottom: 2px solid #eee; padding-bottom: 10px; }}
.summary {{ text-align: center; margin-bottom: 20px; font-size: 1.2em; }}
.card {{ background: #fff; border-radius: 8px; box-shadow: 0 4px 8px rgba(0,0,0,0.1); margin-bottom: 20px; padding: 20px; }}
.layout-container {{ display: flex; gap: 20px; align-items: flex-start; }}
.layout-box {{ flex: 1; text-align: center; }}
.layout-box img {{ max-width: 400px; max-height: 400px; height: auto; border-radius: 4px; cursor: pointer; }}
.masters-grid {{ flex: 2; display: grid; grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); gap: 15px; }}
.master-box {{ text-align: center; border: 1px solid #ddd; border-radius: 8px; padding: 10px; }}
.master-box img {{ width: 100%; height: 150px; object-fit: cover; border-radius: 4px; cursor: pointer; }}
.filename {{ white-space: nowrap; overflow: hidden; text-overflow: ellipsis; font-weight: bold; }}
.alternatives {{ font-size: 0.8em; color: #666; margin-top: 5px; }}
.unmatched-grid {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); gap: 15px; }}
.unmatched-card img {{ width: 100%; height: 150px; object-fit: cover; }}
.modal {{ display: none; position: fixed; z-index: 1000; padding-top: 60px; left: 0; top: 0; width: 100%; height: 100%; overflow: auto; background-color: rgba(0,0,0,0.9); }}
.modal-content {{ margin: auto; display: block; max-width: 90%; max-height: 90vh; }}
.close {{ position: absolute; top: 15px; right: 35px; color: #f1f1f1; font-size: 40px; font-weight: bold; cursor: pointer; }}
</style>
</head><body><div class="container"><h1>Image Match Report</h1><div class="summary">Found matches for {total_matches} out of {total_layouts} layout images.</div>"""
html += "<h2>Matched Layouts</h2>"
for item in data:
if item['found']:
layout_img_path = os.path.join(layouts_abs_path, item['layout']).replace('\\', '/')
enhancements_str = f"<p class='enhancements'>Enhancements: {', '.join(item['enhancements'])}</p>" if item['enhancements'] else ""
html += f"<div class='card'><div class='layout-container'>"
html += f"<div class='layout-box'><h3>Layout Image</h3><img src='file:///{layout_img_path}' onclick='openModal(this.src)'><p class='filename' title='{item['layout']}'>{item['layout']}</p>{enhancements_str}</div>"
html += "<div class='masters-grid'>"
for master_item in item['masters']:
master_img_path = os.path.join(masters_abs_path, master_item['master']).replace('\\', '/')
html += f"<div class='master-box'><img src='file:///{master_img_path}' onclick='openModal(this.src)'><p class='filename' title='{master_item['master']}'>{master_item['master']}</p><p>({master_item['inliers']} inliers)</p></div>"
html += "</div></div></div>"
html += "<h2>Unmatched Layouts</h2><p style='text-align:center;'>Please review these manually.</p><div class='unmatched-grid'>"
for item in data:
if not item['found']:
layout_img_path = os.path.join(layouts_abs_path, item['layout']).replace('\\', '/')
html += f"<div class='card unmatched-card'><img src='file:///{layout_img_path}' onclick='openModal(this.src)'><p class='filename' title='{item['layout']}'>{item['layout']}</p></div>"
html += """</div></div>
<div id="myModal" class="modal"><span class="close" onclick="closeModal()">&times;</span><img class="modal-content" id="img01"></div>
<script>
var modal = document.getElementById("myModal");
var modalImg = document.getElementById("img01");
function openModal(src) { modal.style.display = "block"; modalImg.src = src; }
function closeModal() { modal.style.display = "none"; }
</script>
</body></html>"""
with open(report_path, 'w', encoding='utf-8') as f:
f.write(html)
self.log(f"HTML report saved to {report_path}")
if __name__ == "__main__":
root = tk.Tk()
app = MasterImageFinderApp(root)
root.mainloop()