import gradio as gr import numpy as np import cv2 from PIL import Image import io class AutoWhiteBalance: """自動白平衡處理類""" def __init__(self): self.methods = { "灰世界算法 (Gray World)": "gray_world", "白斑算法 (White Patch)": "white_patch", "簡單平均 (Simple Average)": "simple_avg", "直方圖拉伸 (Histogram Stretch)": "histogram_stretch" } def pil_to_cv2(self, pil_image): """將PIL圖像轉換為OpenCV格式""" # 轉換為RGB（如果是RGBA則去除alpha通道） if pil_image.mode == 'RGBA': pil_image = pil_image.convert('RGB') elif pil_image.mode != 'RGB': pil_image = pil_image.convert('RGB') # 轉換為numpy array image_np = np.array(pil_image) # 轉換為BGR格式（OpenCV默認） image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR) return image_bgr def cv2_to_pil(self, cv2_image): """將OpenCV格式轉換為PIL圖像""" # 轉換回RGB image_rgb = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) # 轉換為PIL圖像 pil_image = Image.fromarray(image_rgb.astype(np.uint8)) return pil_image def gray_world_algorithm(self, image): """灰世界算法 - 假設圖像的平均顏色應該是灰色""" # 計算每個通道的平均值 avg_b = np.mean(image[:, :, 0]) avg_g = np.mean(image[:, :, 1]) avg_r = np.mean(image[:, :, 2]) # 計算整體平均亮度 avg_gray = (avg_b + avg_g + avg_r) / 3 # 避免除零錯誤 scale_b = avg_gray / max(avg_b, 1e-6) scale_g = avg_gray / max(avg_g, 1e-6) scale_r = avg_gray / max(avg_r, 1e-6) # 應用調整 result = image.astype(np.float32) result[:, :, 0] *= scale_b result[:, :, 1] *= scale_g result[:, :, 2] *= scale_r return result def white_patch_algorithm(self, image): """白斑算法 - 假設圖像中最亮的點應該是白色""" # 找到每個通道的最大值（排除極值） max_b = np.percentile(image[:, :, 0], 99) max_g = np.percentile(image[:, :, 1], 99) max_r = np.percentile(image[:, :, 2], 99) # 計算調整係數 scale_b = 255.0 / max(max_b, 1e-6) scale_g = 255.0 / max(max_g, 1e-6) scale_r = 255.0 / max(max_r, 1e-6) # 應用調整 result = image.astype(np.float32) result[:, :, 0] *= scale_b result[:, :, 1] *= scale_g result[:, :, 2] *= scale_r return result def simple_average_algorithm(self, image): """簡單平均算法 - 讓RGB三通道的平均值相等""" # 計算每個通道的平均值 avg_b = np.mean(image[:, :, 0]) avg_g = np.mean(image[:, :, 1]) avg_r = np.mean(image[:, :, 2]) # 以綠色通道為基準（人眼對綠色最敏感） reference = avg_g # 計算調整係數 scale_b = reference / max(avg_b, 1e-6) scale_g = 1.0 # 綠色通道不變 scale_r = reference / max(avg_r, 1e-6) # 應用調整 result = image.astype(np.float32) result[:, :, 0] *= scale_b result[:, :, 1] *= scale_g result[:, :, 2] *= scale_r return result def histogram_stretch_algorithm(self, image): """直方圖拉伸算法""" result = image.astype(np.float32) for i in range(3): # 對每個顏色通道 channel = result[:, :, i] # 計算1%和99%百分位數，忽略極值 p1 = np.percentile(channel, 1) p99 = np.percentile(channel, 99) # 拉伸到0-255範圍 if p99 > p1: channel = (channel - p1) * 255.0 / (p99 - p1) result[:, :, i] = np.clip(channel, 0, 255) return result def preserve_image_brightness(self, original, adjusted): """保持原始圖像的整體亮度""" try: # 計算原始圖像的亮度 original_lab = cv2.cvtColor(original.astype(np.uint8), cv2.COLOR_BGR2LAB) original_brightness = np.mean(original_lab[:, :, 0]) # 計算調整後圖像的亮度 adjusted_lab = cv2.cvtColor(np.clip(adjusted, 0, 255).astype(np.uint8), cv2.COLOR_BGR2LAB) adjusted_brightness = np.mean(adjusted_lab[:, :, 0]) # 調整亮度 if adjusted_brightness > 1e-6: brightness_ratio = original_brightness / adjusted_brightness adjusted_lab[:, :, 0] = np.clip(adjusted_lab[:, :, 0] * brightness_ratio, 0, 255) # 轉換回BGR result = cv2.cvtColor(adjusted_lab, cv2.COLOR_LAB2BGR) return result.astype(np.float32) except: pass return adjusted def process_image(self, pil_image, method_name, strength, preserve_brightness, clip_values): """處理圖像的主函數""" if pil_image is None: return None, "請上傳一張圖片" try: # 轉換格式 cv2_image = self.pil_to_cv2(pil_image) original_image = cv2_image.copy() # 獲取算法名稱 method = self.methods.get(method_name, "gray_world") # 選擇算法 if method == "gray_world": adjusted = self.gray_world_algorithm(cv2_image) algorithm_info = "使用灰世界算法，假設圖像的平均顏色應該是中性灰" elif method == "white_patch": adjusted = self.white_patch_algorithm(cv2_image) algorithm_info = "使用白斑算法，假設圖像中最亮的區域應該是白色" elif method == "simple_avg": adjusted = self.simple_average_algorithm(cv2_image) algorithm_info = "使用簡單平均算法，平衡RGB三個通道的平均值" elif method == "histogram_stretch": adjusted = self.histogram_stretch_algorithm(cv2_image) algorithm_info = "使用直方圖拉伸算法，增強圖像對比度" else: adjusted = cv2_image.astype(np.float32) algorithm_info = "未知算法" # 應用強度調整 if strength != 1.0: adjusted = original_image.astype(np.float32) * (1.0 - strength) + adjusted * strength # 保持亮度 if preserve_brightness: adjusted = self.preserve_image_brightness(original_image, adjusted) # 裁剪數值範圍 if clip_values: adjusted = np.clip(adjusted, 0, 255) # 轉換回PIL格式 result_pil = self.cv2_to_pil(adjusted.astype(np.uint8)) # 生成處理信息 info = f"✅ 處理完成！\n算法：{algorithm_info}\n強度：{strength:.1f}\n保持亮度：{'是' if preserve_brightness else '否'}" return result_pil, info except Exception as e: return None, f"❌ 處理出錯：{str(e)}" # 創建白平衡處理器實例 wb_processor = AutoWhiteBalance() def process_white_balance(image, method, strength, preserve_brightness, clip_values): """Gradio接口函數""" return wb_processor.process_image(image, method, strength, preserve_brightness, clip_values) # 創建Gradio界面 def create_interface(): with gr.Blocks( title="🎨 自動白平衡校正工具", theme=gr.themes.Soft(), css=""" .gradio-container { max-width: 1200px !important; } .image-container { max-height: 600px; } """ ) as demo: gr.Markdown(""" # 🎨 給GPT用的自動白平衡校正工具 上傳有色偏問題的圖片，選擇適合的算法自動校正白平衡，讓圖片恢復自然的色彩！ 💡 **使用建議**： - 🟡 **暖色偏（偏黃）**：推薦使用「灰世界算法」 - 🔵 **冷色偏（偏藍）**：推薦使用「簡單平均」或「白斑算法」 - 🌈 **色彩不鮮豔**：推薦使用「直方圖拉伸，常用」 """) with gr.Row(): with gr.Column(scale=1): # 輸入區域 gr.Markdown("### 📤 上傳圖片") input_image = gr.Image( label="選擇要處理的圖片", type="pil", height=400 ) # 參數設置 gr.Markdown("### ⚙️ 調整參數") method_dropdown = gr.Dropdown( choices=list(wb_processor.methods.keys()), value="灰世界算法 (Gray World)", label="白平衡算法", info="選擇適合的白平衡校正算法" ) strength_slider = gr.Slider( minimum=0.0, maximum=2.0, value=1.0, step=0.1, label="調整強度", info="控制校正效果的強度（0=不調整，1=標準，2=強化）" ) preserve_brightness_checkbox = gr.Checkbox( value=True, label="保持原始亮度", info="保持圖片的整體明暗度不變" ) clip_values_checkbox = gr.Checkbox( value=True, label="裁剪數值範圍", info="避免過度調整造成的色彩異常" ) # 處理按鈕 process_btn = gr.Button( "🚀 開始處理", variant="primary", size="lg" ) with gr.Column(scale=1): # 輸出區域 gr.Markdown("### 📥 處理結果") output_image = gr.Image( label="處理後的圖片", type="pil", height=400 ) # 處理信息 process_info = gr.Textbox( label="處理信息", lines=4, interactive=False ) # 下載提示 gr.Markdown(""" ### 💾 保存圖片 右鍵點擊處理後的圖片，選擇「另存圖片為...」即可保存到本地。 """) # 綁定處理函數 process_btn.click( fn=process_white_balance, inputs=[input_image, method_dropdown, strength_slider, preserve_brightness_checkbox, clip_values_checkbox], outputs=[output_image, process_info] ) # 實時預覽（可選） for component in [method_dropdown, strength_slider, preserve_brightness_checkbox, clip_values_checkbox]: component.change( fn=process_white_balance, inputs=[input_image, method_dropdown, strength_slider, preserve_brightness_checkbox, clip_values_checkbox], outputs=[output_image, process_info] ) # 添加說明 gr.Markdown(""" --- ### 📚 算法說明 - **灰世界算法**：適合校正整體色偏，特別是暖色偏（偏黃/橙） - **白斑算法**：適合有明顯白色或亮色區域的圖片 - **簡單平均**：溫和的校正方式，適合輕微色偏 - **直方圖拉伸**：增強對比度，讓色彩更鮮豔 ### 🔧 技術支持 基於OpenCV和PIL開發，支持常見的圖片格式（JPG、PNG、WEBP等） """) return demo # 啟動應用 if __name__ == "__main__": demo = create_interface() demo.launch( server_name="0.0.0.0", server_port=7860, share=True )