#!/usr/bin/env python3
"""
Convert JPEG to lossless PNG and enhance yellow channel
"""
from PIL import Image
import numpy as np

def convert_and_enhance(input_path, output_path="enhanced_for_deda.png"):
    """Convert to PNG and enhance yellow dots"""
    print(f"Converting {input_path} to enhanced PNG...")
    
    img = Image.open(input_path)
    img_array = np.array(img)
    
    if len(img_array.shape) == 3:
        r = img_array[:,:,0].astype(float)
        g = img_array[:,:,1].astype(float)
        b = img_array[:,:,2].astype(float)
        
        # Create enhanced version where yellow is amplified
        # Make yellow dots more visible
        yellow_score = (r + g - 2*b) / 4
        
        # Amplify yellow
        amplification = 5
        r_new = r + yellow_score * amplification
        g_new = g + yellow_score * amplification
        b_new = b - yellow_score * amplification
        
        # Clip to valid range
        r_new = np.clip(r_new, 0, 255)
        g_new = np.clip(g_new, 0, 255)
        b_new = np.clip(b_new, 0, 255)
        
        enhanced = np.stack([r_new, g_new, b_new], axis=2).astype(np.uint8)
        
        # Save as PNG
        Image.fromarray(enhanced).save(output_path, 'PNG', compress_level=0)
        print(f"Saved enhanced PNG to: {output_path}")
        
        # Also save a version with just yellow channel as grayscale
        yellow_only = ((yellow_score - yellow_score.min()) / (yellow_score.max() - yellow_score.min()) * 255).astype(np.uint8)
        Image.fromarray(yellow_only).save("yellow_only.png", 'PNG')
        print(f"Saved yellow-only channel to: yellow_only.png")
        
        # Create high-contrast version
        threshold = np.percentile(yellow_score, 98)
        high_contrast = np.where(yellow_score > threshold, 255, 0).astype(np.uint8)
        high_contrast_rgb = np.stack([high_contrast, high_contrast, np.zeros_like(high_contrast)], axis=2)
        Image.fromarray(high_contrast_rgb).save("high_contrast_dots.png", 'PNG')
        print(f"Saved high-contrast version to: high_contrast_dots.png")

if __name__ == "__main__":
    convert_and_enhance("Zard - Copy.jpg")