3_Best_Example_Selector.py

#!/usr/bin/env python3
"""
HPA ROI Example Selector
Loads ROI.txt, adds annotation labels, and allows selection of example images.
"""

import csv
from pathlib import Path
import tkinter as tk
from tkinter import ttk
from PIL import Image, ImageTk
import math


def load_roi_file(roi_file_path):
    """
    Load ROI data from CSV file.
    
    Args:
        roi_file_path: Path to the ROI.txt file
    
    Returns:
        list: List of dictionaries containing ROI data
    """
    rows = []
    
    if not roi_file_path.exists():
        print(f"Error: {roi_file_path} not found!")
        return rows
    
    with open(roi_file_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        rows = list(reader)
    
    return rows


def load_annotation_file(gene_name, script_dir):
    """
    Load annotation from anno.txt file for a given gene.
    
    Args:
        gene_name: Name of the gene
        script_dir: Script directory
    
    Returns:
        str: Annotation text or 'N/A' if not found
    """
    anno_path = script_dir / gene_name / "anno.txt"
    
    if anno_path.exists():
        try:
            with open(anno_path, 'r', encoding='utf-8') as f:
                return f.read().strip()
        except:
            return 'N/A'
    return 'N/A'


def group_by_gene(roi_data):
    """
    Group ROI data by gene.
    
    Args:
        roi_data: List of ROI dictionaries
    
    Returns:
        dict: Dictionary mapping gene names to lists of ROI entries
    """
    gene_groups = {}
    
    for row in roi_data:
        gene = row['Gene']
        if gene not in gene_groups:
            gene_groups[gene] = []
        gene_groups[gene].append(row)
    
    return gene_groups


class ImageSelector:
    """GUI for selecting example images from a grid."""
    
    def __init__(self, gene_name, annotation, roi_entries, script_dir, thumbnail_size=200):
        """
        Initialize the image selector.
        
        Args:
            gene_name: Name of the gene
            annotation: Annotation label
            roi_entries: List of ROI entries for this gene
            script_dir: Script directory
            thumbnail_size: Size of thumbnail images
        """
        self.gene_name = gene_name
        self.annotation = annotation
        self.roi_entries = roi_entries
        self.script_dir = script_dir
        self.thumbnail_size = thumbnail_size
        self.selected_index = None
        self.selected_buttons = []
        
        # Create GUI
        self.setup_gui()
    
    def setup_gui(self):
        """Setup the GUI window."""
        self.root = tk.Tk()
        self.root.title(f"Select Example Image - {self.gene_name}")
        
        # Main frame
        main_frame = tk.Frame(self.root, padx=20, pady=20)
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # Header
        header_text = f"Gene: {self.gene_name}\nAnnotation: {self.annotation}\n\nSelect ONE example image:"
        header_label = tk.Label(main_frame, text=header_text, 
                               font=('Arial', 14, 'bold'), justify=tk.LEFT)
        header_label.pack(pady=(0, 20))
        
        # Create scrollable frame for images
        canvas_frame = tk.Frame(main_frame)
        canvas_frame.pack(fill=tk.BOTH, expand=True)
        
        # Canvas with scrollbar
        canvas = tk.Canvas(canvas_frame, bg='white')
        scrollbar = ttk.Scrollbar(canvas_frame, orient="vertical", command=canvas.yview)
        scrollable_frame = tk.Frame(canvas, bg='white')
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
        # Calculate grid dimensions (max 4 columns)
        num_images = len(self.roi_entries)
        num_cols = min(4, num_images)
        num_rows = math.ceil(num_images / num_cols)
        
        # Create image grid
        self.images = []
        self.photo_images = []
        
        for idx, entry in enumerate(self.roi_entries):
            row_idx = idx // num_cols
            col_idx = idx % num_cols
            
            # Create frame for each image
            img_frame = tk.Frame(scrollable_frame, relief=tk.RAISED, borderwidth=2, bg='white')
            img_frame.grid(row=row_idx, column=col_idx, padx=10, pady=10, sticky='nsew')
            
            # Load and display image
            try:
                image_path = self.get_image_path(entry)
                
                if image_path and image_path.exists():
                    # Load image with ROI overlay
                    img = self.create_thumbnail_with_roi(image_path, entry)
                    photo = ImageTk.PhotoImage(img)
                    self.photo_images.append(photo)  # Keep reference
                    
                    # Create button with image
                    btn = tk.Button(img_frame, image=photo, 
                                   command=lambda i=idx: self.select_image(i),
                                   relief=tk.FLAT, bg='white')
                    btn.pack(padx=5, pady=5)
                    self.selected_buttons.append(btn)
                    
                    # Image info
                    info_text = f"{entry['ImagePrefix']}\n{entry['Antibody']}\n{entry['CellLine']}"
                    if entry['Skip'] == 'T':
                        info_text += "\n[SKIP]"
                    
                    info_label = tk.Label(img_frame, text=info_text, 
                                         font=('Arial', 9), justify=tk.CENTER, bg='white')
                    info_label.pack()
                else:
                    # Image not found
                    error_label = tk.Label(img_frame, text="Image not found", 
                                          font=('Arial', 10), fg='red', bg='white')
                    error_label.pack(pady=50)
                    
            except Exception as e:
                print(f"Error loading image {idx}: {e}")
                error_label = tk.Label(img_frame, text=f"Error: {str(e)}", 
                                      font=('Arial', 9), fg='red', bg='white')
                error_label.pack(pady=50)
        
        # Button frame
        button_frame = tk.Frame(main_frame)
        button_frame.pack(pady=(20, 0))
        
        # OK button
        ok_button = tk.Button(button_frame, text="Confirm Selection", 
                             command=self.on_ok,
                             bg="green", fg="white", padx=30, pady=10,
                             font=('Arial', 12, 'bold'))
        ok_button.pack(side=tk.LEFT, padx=5)
        
        # Skip All button
        skip_button = tk.Button(button_frame, text="Skip All (No Example)", 
                               command=self.on_skip_all,
                               bg="orange", fg="white", padx=30, pady=10,
                               font=('Arial', 12, 'bold'))
        skip_button.pack(side=tk.LEFT, padx=5)
        
        # Bind mouse wheel for scrolling
        canvas.bind_all("<MouseWheel>", lambda e: canvas.yview_scroll(int(-1*(e.delta/120)), "units"))
        
        # Set window size
        window_width = min(1400, num_cols * (self.thumbnail_size + 60))
        window_height = min(900, num_rows * (self.thumbnail_size + 100) + 300)
        
        # Center window
        self.root.update_idletasks()
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()
        x = (screen_width // 2) - (window_width // 2)
        y = (screen_height // 2) - (window_height // 2)
        self.root.geometry(f"{window_width}x{window_height}+{x}+{y}")
    
    def get_image_path(self, entry):
        """Get the full path to the composite image."""
        folder_path = self.script_dir / entry['FolderPath']
        image_name = f"{entry['ImagePrefix']}_blue_red_green.jpg"
        return folder_path / image_name
    
    def create_thumbnail_with_roi(self, image_path, entry):
        """Create a thumbnail with ROI overlay."""
        from PIL import ImageDraw
        
        # Load image
        img = Image.open(image_path)
        
        # Draw ROI rectangle
        draw = ImageDraw.Draw(img)
        roi_x = int(entry['Roi_X'])
        roi_y = int(entry['Roi_Y'])
        roi_size = 640
        
        draw.rectangle(
            [roi_x, roi_y, roi_x + roi_size, roi_y + roi_size],
            outline="red", width=8
        )
        
        # Draw crosshair
        center_x = roi_x + roi_size // 2
        center_y = roi_y + roi_size // 2
        crosshair_size = 30
        draw.line([center_x - crosshair_size, center_y, 
                   center_x + crosshair_size, center_y], fill="red", width=6)
        draw.line([center_x, center_y - crosshair_size, 
                   center_x, center_y + crosshair_size], fill="red", width=6)
        
        # Resize to thumbnail
        img.thumbnail((self.thumbnail_size, self.thumbnail_size), Image.Resampling.LANCZOS)
        
        return img
    
    def select_image(self, index):
        """Handle image selection."""
        # Reset all buttons to normal
        for btn in self.selected_buttons:
            btn.config(relief=tk.FLAT, bg='white', borderwidth=0)
        
        # Highlight selected button
        self.selected_buttons[index].config(relief=tk.SUNKEN, bg='lightgreen', borderwidth=5)
        self.selected_index = index
    
    def on_ok(self):
        """Handle OK button click."""
        if self.selected_index is None:
            tk.messagebox.showwarning("No Selection", 
                                     "Please select an image or click 'Skip All'")
            return
        
        self.root.quit()
        self.root.destroy()
    
    def on_skip_all(self):
        """Handle Skip All button click."""
        self.selected_index = -1  # Special value to indicate skip all
        self.root.quit()
        self.root.destroy()
    
    def run(self):
        """Run the selector and return the selected index."""
        self.root.mainloop()
        return self.selected_index


def save_roi_examples(output_path, roi_data_with_examples):
    """
    Save ROI data with annotation and example columns to CSV.
    
    Args:
        output_path: Path to save the file
        roi_data_with_examples: List of ROI dictionaries with added columns
    """
    fieldnames = ['Gene', 'Antibody', 'CellLine', 'ImagePrefix', 'FolderPath', 
                  'Roi_X', 'Roi_Y', 'Skip', 'Annotation', 'Example']
    
    with open(output_path, 'w', encoding='utf-8', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(roi_data_with_examples)
    
    print(f"\nSaved to: {output_path}")


def main():
    """Main function."""
    print("=" * 70)
    print("HPA ROI Example Selector")
    print("=" * 70)
    
    # Get script directory
    script_dir = Path(__file__).parent
    
    # Input and output file paths
    roi_file_path = script_dir / "ROI.txt"
    output_file_path = script_dir / "ROI_examples.txt"
    
    print(f"\nInput file: {roi_file_path}")
    print(f"Output file: {output_file_path}")
    
    # Ask user for selection mode
    print("\n" + "=" * 70)
    print("Select example selection mode:")
    print("=" * 70)
    print("1. Complement only - select examples for genes without selections")
    print("2. Re-select all - allow re-selection for all genes")
    print("=" * 70)
    
    while True:
        mode = input("\nEnter mode (1 or 2): ").strip()
        if mode in ['1', '2']:
            break
        print("Invalid input. Please enter 1 or 2.")
    
    complement_only = (mode == '1')
    
    if complement_only:
        print("\nMode: Complement only - will skip genes with existing example selections")
    else:
        print("\nMode: Re-select all - will allow re-selection for all genes")
    
    # Load ROI data
    print("\nLoading ROI data...")
    roi_data = load_roi_file(roi_file_path)
    
    if not roi_data:
        print("No ROI data found!")
        return
    
    print(f"Loaded {len(roi_data)} ROI entries")
    
    # Check if ROI_examples.txt exists and load existing examples
    existing_examples = {}
    if output_file_path.exists():
        print(f"\nLoading existing examples from {output_file_path}...")
        try:
            with open(output_file_path, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    if row.get('Example') == 'T':
                        gene = row['Gene']
                        existing_examples[gene] = {
                            'Antibody': row['Antibody'],
                            'CellLine': row['CellLine'],
                            'ImagePrefix': row['ImagePrefix']
                        }
            print(f"Found {len(existing_examples)} existing example selections")
        except Exception as e:
            print(f"Warning: Could not load existing examples: {e}")
            existing_examples = {}
    
    # Group by gene
    print("\nGrouping by gene...")
    gene_groups = group_by_gene(roi_data)
    print(f"Found {len(gene_groups)} unique genes")
    
    # Add annotations to all entries
    print("\nAdding annotations...")
    for row in roi_data:
        gene = row['Gene']
        annotation = load_annotation_file(gene, script_dir)
        row['Annotation'] = annotation
        row['Example'] = 'F'  # Default to False
    
    # Filter genes based on mode
    if complement_only:
        genes_to_process = {gene: entries for gene, entries in gene_groups.items() 
                           if gene not in existing_examples}
        print(f"\nGenes without examples: {len(genes_to_process)}")
        
        if not genes_to_process:
            print("All genes already have examples selected!")
            
            # Still apply existing examples to the data
            for row in roi_data:
                gene = row['Gene']
                if gene in existing_examples:
                    ex = existing_examples[gene]
                    if (row['Antibody'] == ex['Antibody'] and 
                        row['CellLine'] == ex['CellLine'] and 
                        row['ImagePrefix'] == ex['ImagePrefix']):
                        row['Example'] = 'T'
            
            # Save with existing examples
            save_roi_examples(output_file_path, roi_data)
            return
    else:
        genes_to_process = gene_groups
        print(f"\nWill process all {len(genes_to_process)} genes")
    
    # Apply existing examples first (they will be overwritten in re-select mode if user chooses)
    for row in roi_data:
        gene = row['Gene']
        if gene in existing_examples:
            ex = existing_examples[gene]
            if (row['Antibody'] == ex['Antibody'] and 
                row['CellLine'] == ex['CellLine'] and 
                row['ImagePrefix'] == ex['ImagePrefix']):
                row['Example'] = 'T'
    
    # Process each gene
    total_genes = len(genes_to_process)
    processed = 0
    
    for gene_idx, (gene_name, gene_entries) in enumerate(sorted(genes_to_process.items()), 1):
        print("\n" + "=" * 70)
        print(f"Gene {gene_idx}/{total_genes}: {gene_name}")
        print("=" * 70)
        
        annotation = gene_entries[0]['Annotation']  # All entries have same annotation
        print(f"Annotation: {annotation}")
        print(f"Number of images: {len(gene_entries)}")
        
        # Check if example exists and show info
        if gene_name in existing_examples:
            ex = existing_examples[gene_name]
            print(f"Existing example: {ex['ImagePrefix']} ({ex['Antibody']}/{ex['CellLine']})")
            
            if not complement_only:
                response = input("Re-select example for this gene? (y/n, default=n): ").strip().lower()
                if response != 'y':
                    print("Keeping existing selection...\n")
                    continue
        
        # Show image selector
        try:
            selector = ImageSelector(gene_name, annotation, gene_entries, script_dir)
            selected_idx = selector.run()
            
            # First, clear any existing examples for this gene
            for entry in gene_entries:
                entry['Example'] = 'F'
            
            if selected_idx == -1:
                print("No example selected for this gene (skipped)")
            elif selected_idx is not None:
                # Mark the selected image as example
                selected_entry = gene_entries[selected_idx]
                selected_entry['Example'] = 'T'
                print(f"Selected: {selected_entry['ImagePrefix']} ({selected_entry['Antibody']}/{selected_entry['CellLine']})")
                processed += 1
            else:
                print("Selection cancelled")
        
        except Exception as e:
            print(f"Error processing gene {gene_name}: {e}")
            import traceback
            traceback.print_exc()
    
    # Save results
    print("\n" + "=" * 70)
    print("Saving results...")
    print("=" * 70)
    
    save_roi_examples(output_file_path, roi_data)
    
    print(f"\nProcessed {processed} genes with new selections")
    print(f"Total entries: {len(roi_data)}")
    
    # Summary
    examples_count = sum(1 for row in roi_data if row['Example'] == 'T')
    print(f"Total examples selected: {examples_count}")
    
    print("\n" + "=" * 70)
    print("Complete!")
    print("=" * 70)


if __name__ == "__main__":
    main()