inference.py

import tensorflow as tf
import cv2
import numpy as np
from PIL import Image
import argparse
import os

class MangaColorizerInference:
    def __init__(self, model_path):
        """โหลดโมเดลที่ฝึกไว้แล้ว"""
        try:
            if not os.path.exists(model_path):
                raise FileNotFoundError(f"ไม่พบไฟล์โมเดล: {model_path}")
            
            self.model = tf.keras.models.load_model(model_path)
            print(f"✅ โหลดโมเดลสำเร็จจาก: {model_path}")
            
        except Exception as e:
            print(f"❌ เกิดข้อผิดพลาดในการโหลดโมเดล: {e}")
            raise
    
    def load_and_preprocess(self, image_path):
        """โหลดและเตรียมภาพสำหรับโมเดล"""
        try:
            # โหลดภาพ
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"ไม่สามารถโหลดภาพจาก: {image_path}")
            
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # บันทึกขนาดเดิม
            original_size = image.shape[:2]
            
            # ปรับขนาด
            image_resized = cv2.resize(image, (256, 256))
            
            # แปลงเป็น grayscale
            gray = cv2.cvtColor(image_resized, cv2.COLOR_RGB2GRAY)
            
            # ทำให้เส้นคมชัด
            _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)
            binary = cv2.bitwise_not(binary)
            
            # Normalize
            binary = binary.astype(np.float32) / 255.0
            binary = np.expand_dims(binary, axis=-1)
            
            return binary, original_size, image_resized
            
        except Exception as e:
            print(f"❌ ข้อผิดพลาดในการประมวลผลภาพ: {e}")
            raise
    
    def colorize(self, image_path, output_path, save_comparison=False):
        """ลงสีภาพเส้น"""
        try:
            # โหลดและประมวลผลภาพ
            image, original_size, original_resized = self.load_and_preprocess(image_path)
            
            # ทำนายสี
            print("🎨 กำลังลงสี...")
            prediction = self.model.predict(image[np.newaxis, ...], verbose=0)[0]
            
            # แปลงกลับเป็นภาพและปรับขนาดกลับ
            colored_image = (prediction * 255).astype(np.uint8)
            colored_image_original_size = cv2.resize(colored_image, (original_size[1], original_size[0]))
            
            # บันทึกผลลัพธ์
            Image.fromarray(colored_image_original_size).save(output_path)
            print(f"💾 บันทึกภาพที่: {output_path}")
            
            # บันทึกภาพเปรียบเทียบ (optional)
            if save_comparison:
                comparison = self.create_comparison(original_resized, colored_image)
                comparison_path = output_path.replace('.png', '_comparison.png')
                comparison.save(comparison_path)
                print(f"📊 บันทึกภาพเปรียบเทียบที่: {comparison_path}")
            
            return colored_image_original_size
            
        except Exception as e:
            print(f"❌ ข้อผิดพลาดในการลงสี: {e}")
            return None
    
    def create_comparison(self, original, colored):
        """สร้างภาพเปรียบเทียบก่อน-หลัง"""
        # แปลงภาพต้นฉบับเป็น RGB สำหรับแสดง
        if len(original.shape) == 3 and original.shape[2] == 3:
            original_rgb = original
        else:
            original_rgb = cv2.cvtColor(original, cv2.COLOR_GRAY2RGB)
        
        # รวมภาพ
        comparison = np.hstack([original_rgb, colored])
        return Image.fromarray(comparison)

def main():
    parser = argparse.ArgumentParser(description='AI Manga Colorizer - Inference')
    parser.add_argument('--input', type=str, required=True, help='Path to input line art')
    parser.add_argument('--output', type=str, required=True, help='Path to save colored image')
    parser.add_argument('--model', type=str, default='output/manga_colorizer.h5', help='Path to trained model')
    parser.add_argument('--compare', action='store_true', help='Save comparison image')
    
    args = parser.parse_args()
    
    # ตรวจสอบไฟล์ input
    if not os.path.exists(args.input):
        print(f"❌ ไม่พบไฟล์ input: {args.input}")
        return
    
    # ใช้งานโมเดล
    try:
        colorizer = MangaColorizerInference(args.model)
        result = colorizer.colorize(args.input, args.output, save_comparison=args.compare)
        
        if result is not None:
            print("✅ เสร็จสิ้น!")
        else:
            print("❌ การลงสีล้มเหลว")
        
    except Exception as e:
        print(f"❌ การทำงานล้มเหลว: {e}")

if __name__ == "__main__":
    main()