app.py

import gradio as gr
from utils import WatermarkProcessor
import json
import tempfile
import os
from datetime import datetime
import cv2
from PIL import Image
import numpy as np
import logging

# 로깅 설정 (디버그 정보 출력)
logging.basicConfig(level=logging.DEBUG)

class WatermarkGUI:
    def __init__(self):
        self.processor = WatermarkProcessor()
        self.create_interface()

    def process_watermark(self, image, watermark_text, author, purpose, opacity):
        """메타데이터와 함께 워터마크를 이미지에 추가"""
        if image is None or watermark_text.strip() == "":
            return None, "이미지와 워터마크 텍스트를 모두 제공해주세요."
        
        metadata = {
            "author": author,
            "purpose": purpose,
            "opacity": opacity
        }
        
        # 임시 이미지 저장
        temp_path = tempfile.mktemp(suffix='.png')
        try:
            Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).save(temp_path)
            logging.debug("임시 이미지 저장 완료: %s", temp_path)
        except Exception as e:
            logging.error("임시 이미지 저장 중 오류: %s", e)
            return None, f"이미지 저장 오류: {e}"
        
        # 워터마크 추가
        result_path, message = self.processor.encode(temp_path, watermark_text, metadata)
        logging.debug("워터마크 인코딩 결과 - result_path: %s, message: %s", result_path, message)
        
        if "Error" in message:
            os.remove(temp_path)
            return None, message
        
        # 품질 보고서 생성
        quality_report = self.processor.analyze_quality(temp_path, result_path)
        try:
            quality_data = json.loads(quality_report)
        except Exception as e:
            logging.error("품질 보고서 파싱 오류: %s", e)
            quality_data = {"quality_score": "N/A", "psnr": "N/A", "histogram_similarity": 0, "modified_pixels": 0}
        
        report = f"""
### Watermark Quality Report
- Quality Score: {quality_data.get('quality_score', 'N/A')}%
- PSNR: {quality_data.get('psnr', 'N/A')} dB
- Histogram Similarity: {quality_data.get('histogram_similarity', 0) * 100:.2f}%
- Modified Pixels: {quality_data.get('modified_pixels', 0):,}

### Metadata
- Author: {author}
- Purpose: {purpose}
- Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
        """
        
        os.remove(temp_path)
        watermarked_image = cv2.imread(result_path)
        if watermarked_image is None:
            logging.error("워터마크가 추가된 이미지 로드 실패: %s", result_path)
            return None, "워터마크 이미지 로드에 실패했습니다."
        return watermarked_image, report

    def detect_watermark(self, image):
        """워터마크 검출 및 추출"""
        if image is None:
            return "이미지를 제공해주세요."
        
        # 임시 이미지 저장
        temp_path = tempfile.mktemp(suffix='.png')
        try:
            Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).save(temp_path)
            logging.debug("워터마크 검출용 임시 이미지 저장 완료: %s", temp_path)
        except Exception as e:
            logging.error("검출용 이미지 저장 중 오류: %s", e)
            return f"이미지 저장 오류: {e}"
        
        # 워터마크 추출
        result = self.processor.decode(temp_path)
        logging.debug("디코드 결과: %s", result)
        os.remove(temp_path)
        
        try:
            # JSON 파싱 시도
            data = json.loads(result)
            report = f"""
### Extracted Watermark
Text: {data.get('text', 'N/A')}

### Metadata
- Timestamp: {data.get('timestamp', 'N/A')}
- Author: {data.get('metadata', {}).get('author', 'N/A')}
- Purpose: {data.get('metadata', {}).get('purpose', 'N/A')}
            """
            return report
        except Exception as e:
            logging.error("워터마크 디코딩 결과 파싱 오류: %s", e)
            # JSON 파싱 실패 시, 원본 결과 반환
            return result

    def create_interface(self):
        """Gradio 인터페이스 생성"""
        with gr.Blocks(css="footer {visibility: hidden}") as self.interface:
            gr.HTML(
                """<a href="https://visitorbadge.io/status?path=https%3A%2F%2Ffantos-SecureWatermark.hf.space">
                <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Ffantos-SecureWatermark.hf.space&countColor=%23263759" />
                </a>"""
            )
            
            gr.Markdown(
                """# Enhanced Image Watermarking System
                         
### Welcome to Secure Watermark - Advanced Image Protection System

🔒 **Key Features:**
- **Dual Watermarking Technology**: Supports both steganography and PNG metadata
- **Secure Encryption**: Military-grade encryption for watermark data
- **Quality Assurance**: Real-time quality analysis and reporting
- **Metadata Support**: Track authorship, purpose, and timestamps
- **Integrity Verification**: Hash-based image tampering detection

💡 **Perfect for:**
- Copyright Protection
- Digital Asset Management
- Document Authentication
- Creative Work Protection

Try our system by uploading an image and adding your watermark below!
                """
            )
            
            with gr.Tabs():
                # 워터마크 추가 탭
                with gr.Tab("Add Watermark"):
                    with gr.Row():
                        with gr.Column():
                            input_image = gr.Image(label="Input Image", type="numpy")
                            watermark_text = gr.Textbox(label="Watermark Text")
                            author = gr.Textbox(label="Author", placeholder="Enter author name")
                            purpose = gr.Textbox(label="Purpose", placeholder="Enter watermark purpose")
                            opacity = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, label="Watermark Opacity")
                            
                            with gr.Row():
                                process_btn = gr.Button("Add Watermark", variant="primary")
                        with gr.Column():
                            result_image = gr.Image(label="Watermarked Image")
                            quality_report = gr.Markdown(label="Quality Report")
                
                # 워터마크 검출 탭
                with gr.Tab("Detect Watermark"):
                    with gr.Row():
                        detect_image = gr.Image(label="Input Image", type="numpy")
                        detect_result = gr.Markdown(label="Detected Watermark")
                    detect_btn = gr.Button("Detect Watermark")
            
            # 이벤트 핸들러 연결
            process_btn.click(
                fn=self.process_watermark,
                inputs=[input_image, watermark_text, author, purpose, opacity],
                outputs=[result_image, quality_report]
            )
            
            detect_btn.click(
                fn=self.detect_watermark,
                inputs=[detect_image],
                outputs=detect_result
            )

    def launch(self, *args, **kwargs):
        """인터페이스 실행"""
        self.interface.launch(*args, **kwargs)

if __name__ == "__main__":
    app = WatermarkGUI()
    app.launch()