devmeta's picture
Upload 2 files
48701bd verified
# -*- coding: utf-8 -*-
"""워터마킹시스템.ipynb
Automatically generated by Colab.
Original file is located at
https://colab.research.google.com/drive/17EOpDL6hwJ6f3G_-7bpUspm-1-XcgL_w
"""
import gradio as gr
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import io
import base64
import json
import time
from typing import Tuple, Optional
import matplotlib.pyplot as plt
import tempfile
import os
# ===== 경량화 CNN 워터마킹 모델 =====
class MobileWatermarkEncoder(nn.Module):
"""모바일 최적화 워터마크 인코더"""
def __init__(self, watermark_size=32):
super().__init__()
self.watermark_size = watermark_size
# 경량화 인코더 (MobileNet 스타일)
self.encoder = nn.Sequential(
# 초기 특징 추출
nn.Conv2d(3, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
# Depthwise Separable Convolution
nn.Conv2d(32, 32, 3, padding=1, groups=32), # Depthwise
nn.Conv2d(32, 64, 1), # Pointwise
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, 3, padding=1, groups=64),
nn.Conv2d(64, 128, 1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
# 출력층
nn.Conv2d(128, 3, 3, padding=1),
nn.Tanh()
)
# 워터마크 임베딩 강도 조절
self.alpha = nn.Parameter(torch.tensor(0.1))
def forward(self, image, watermark_pattern):
# 워터마크 패턴을 이미지 크기에 맞게 확장
h, w = image.shape[2], image.shape[3]
watermark = F.interpolate(watermark_pattern, size=(h, w), mode='bilinear')
# 워터마크 임베딩
watermark_noise = self.encoder(image)
watermarked = image + self.alpha * watermark_noise * watermark
return torch.clamp(watermarked, 0, 1)
class MobileWatermarkDecoder(nn.Module):
"""모바일 최적화 워터마크 디코더"""
def __init__(self, watermark_size=32):
super().__init__()
self.watermark_size = watermark_size
self.decoder = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d((watermark_size, watermark_size)),
nn.Conv2d(64, 1, 1),
nn.Sigmoid()
)
def forward(self, image):
return self.decoder(image)
# ===== 워터마킹 시스템 클래스 =====
class MobileWatermarkingSystem:
def __init__(self):
self.encoder = MobileWatermarkEncoder()
self.decoder = MobileWatermarkDecoder()
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 모델을 device로 이동
self.encoder.to(self.device)
self.decoder.to(self.device)
# 간단한 훈련용 더미 데이터로 초기화
self._initialize_models()
def _initialize_models(self):
"""모델 초기화 (실제로는 사전 훈련된 가중치 로드)"""
# 여기서는 간단한 초기화만 수행
# 실제 구현에서는 사전 훈련된 모델 로드
pass
def generate_watermark_pattern(self, user_id: str, timestamp: str) -> torch.Tensor:
"""사용자 ID와 타임스탬프로 고유 워터마크 패턴 생성"""
# 간단한 패턴 생성 (실제로는 더 복잡한 방법 사용)
seed = hash(user_id + timestamp) % 10000
torch.manual_seed(seed)
pattern = torch.randn(1, 1, 32, 32)
return torch.sigmoid(pattern)
def embed_watermark(self, image: np.ndarray, user_id: str) -> Tuple[np.ndarray, dict]:
"""이미지에 워터마크 삽입"""
start_time = time.time()
# 이미지 전처리
if len(image.shape) == 3:
image_tensor = torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0
else:
image_tensor = torch.from_numpy(image).float() / 255.0
image_tensor = image_tensor.unsqueeze(0).repeat(3, 1, 1)
image_tensor = image_tensor.unsqueeze(0).to(self.device)
# 워터마크 패턴 생성
timestamp = str(int(time.time()))
watermark_pattern = self.generate_watermark_pattern(user_id, timestamp)
watermark_pattern = watermark_pattern.to(self.device)
# 워터마크 삽입
with torch.no_grad():
watermarked_tensor = self.encoder(image_tensor, watermark_pattern)
# 후처리
watermarked_image = watermarked_tensor.squeeze(0).cpu().numpy()
watermarked_image = (watermarked_image.transpose(1, 2, 0) * 255).astype(np.uint8)
processing_time = time.time() - start_time
# 메타데이터
metadata = {
'user_id': user_id,
'timestamp': timestamp,
'processing_time': processing_time,
'image_size': image.shape,
'watermark_strength': float(self.encoder.alpha.item())
}
return watermarked_image, metadata
def extract_watermark(self, image: np.ndarray) -> Tuple[np.ndarray, float]:
"""이미지에서 워터마크 추출 및 검증"""
start_time = time.time()
# 이미지 전처리
if len(image.shape) == 3:
image_tensor = torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0
else:
image_tensor = torch.from_numpy(image).float() / 255.0
image_tensor = image_tensor.unsqueeze(0).repeat(3, 1, 1)
image_tensor = image_tensor.unsqueeze(0).to(self.device)
# 워터마크 추출
with torch.no_grad():
extracted_watermark = self.decoder(image_tensor)
# 후처리
watermark_array = extracted_watermark.squeeze().cpu().numpy()
confidence = np.mean(watermark_array) # 간단한 신뢰도 계산
processing_time = time.time() - start_time
return watermark_array, confidence
def verify_watermark(self, original_metadata: dict, extracted_confidence: float) -> dict:
"""워터마크 검증"""
threshold = 0.3 # 검증 임계값
is_valid = bool(extracted_confidence > threshold) # bool() 명시적 변환
return {
'is_valid': is_valid,
'confidence': float(extracted_confidence), # float() 명시적 변환
'threshold': float(threshold),
'original_metadata': original_metadata
}
# ===== 전역 시스템 인스턴스 =====
watermarking_system = MobileWatermarkingSystem()
# ===== Gradio 인터페이스 함수들 =====
def embed_watermark_interface(image, user_id, output_format):
"""워터마크 삽입 인터페이스"""
if image is None:
return None, "이미지를 업로드해주세요.", None, None
if not user_id.strip():
return None, "사용자 ID를 입력해주세요.", None, None
try:
# 워터마크 삽입
watermarked_image, metadata = watermarking_system.embed_watermark(image, user_id)
# 메타데이터를 JSON으로 변환
metadata_json = json.dumps(metadata, indent=2)
# 다운로드용 파일 생성
download_file = create_download_file(watermarked_image, user_id, metadata, output_format)
# 결과 메시지
result_msg = f"""
✅ 워터마크 삽입 완료!
📊 처리 시간: {metadata['processing_time']:.3f}
📏 이미지 크기: {metadata['image_size']}
💪 워터마크 강도: {metadata['watermark_strength']:.3f}
🆔 사용자 ID: {metadata['user_id']}
⏰ 타임스탬프: {metadata['timestamp']}
📄 포맷: {output_format.upper()}
"""
return watermarked_image, result_msg, metadata_json, download_file
except Exception as e:
return None, f"오류 발생: {str(e)}", None, None
def create_download_file(image, user_id, metadata, output_format):
"""다운로드용 파일 생성"""
try:
# PIL Image로 변환
if isinstance(image, np.ndarray):
pil_image = Image.fromarray(image)
else:
pil_image = image
# 파일명 생성
timestamp = metadata['timestamp']
filename = f"watermarked_{user_id}_{timestamp}.{output_format.lower()}"
# 임시 파일 생성
temp_dir = tempfile.mkdtemp()
temp_path = os.path.join(temp_dir, filename)
if output_format.lower() == 'jpg':
# JPG는 RGB 모드 필요
if pil_image.mode in ('RGBA', 'LA', 'P'):
# 투명도가 있는 경우 흰색 배경과 합성
background = Image.new('RGB', pil_image.size, (255, 255, 255))
if pil_image.mode == 'P':
pil_image = pil_image.convert('RGBA')
background.paste(pil_image, mask=pil_image.split()[-1] if pil_image.mode == 'RGBA' else None)
pil_image = background
pil_image.save(temp_path, format='JPEG', quality=95, optimize=True)
elif output_format.lower() == 'png':
# PNG 메타데이터
pnginfo = Image.PngImagePlugin.PngInfo()
pnginfo.add_text("User_ID", user_id)
pnginfo.add_text("Timestamp", timestamp)
pnginfo.add_text("Watermark_Strength", str(metadata['watermark_strength']))
pnginfo.add_text("Software", "Mobile Watermarking System")
pil_image.save(temp_path, format='PNG', pnginfo=pnginfo, optimize=True)
return temp_path
except Exception as e:
print(f"다운로드 파일 생성 오류: {e}")
return None
def extract_watermark_interface(image, metadata_json):
"""워터마크 추출 및 검증 인터페이스"""
if image is None:
return None, "이미지를 업로드해주세요.", None
try:
# 워터마크 추출
watermark_pattern, confidence = watermarking_system.extract_watermark(image)
# 워터마크 패턴 시각화 (제목만 영어, 폰트 설정)
plt.figure(figsize=(6, 6))
plt.rcParams['font.family'] = 'DejaVu Sans' # 영어 폰트 설정
plt.rcParams['font.size'] = 10
plt.imshow(watermark_pattern, cmap='viridis')
plt.title(f'Extracted Watermark Pattern (Confidence: {confidence:.3f})',
fontsize=12, fontweight='bold')
plt.colorbar(label='Pattern Intensity')
plt.axis('off')
# 임시 파일로 저장
watermark_viz = plt.gcf()
# 신뢰도별 해석 메시지 생성 (한글)
def get_confidence_interpretation(conf):
if conf >= 0.8:
return {
'level': '높음',
'emoji': '✅',
'message': '워터마크가 명확히 감지되었습니다.',
'detail': '이 이미지는 워터마크가 삽입된 이미지로 판단됩니다.',
'color': '🟢'
}
elif conf >= 0.6:
return {
'level': '보통',
'emoji': '⚠️',
'message': '워터마크 패턴이 감지되었습니다.',
'detail': '워터마크가 있을 가능성이 높지만 추가 검증이 권장됩니다.',
'color': '🟡'
}
elif conf >= 0.3:
return {
'level': '낮음',
'emoji': '❓',
'message': '약한 워터마크 신호가 감지되었습니다.',
'detail': '워터마크가 있을 수 있지만 노이즈일 가능성도 있습니다.',
'color': '🟠'
}
else:
return {
'level': '매우 낮음',
'emoji': '❌',
'message': '워터마크가 감지되지 않았습니다.',
'detail': '이 이미지에는 워터마크가 없거나 손상되었을 가능성이 높습니다.',
'color': '🔴'
}
confidence_info = get_confidence_interpretation(confidence)
# 검증 수행
verification_result = None
if metadata_json and metadata_json.strip():
try:
original_metadata = json.loads(metadata_json)
verification_result = watermarking_system.verify_watermark(
original_metadata, confidence
)
# JSON 직렬화 가능하도록 데이터 타입 보장
verification_result = {
'is_valid': bool(verification_result['is_valid']),
'confidence': float(verification_result['confidence']),
'threshold': float(verification_result['threshold']),
'original_metadata': verification_result['original_metadata'],
'confidence_level': confidence_info['level'],
'interpretation': confidence_info['message']
}
except json.JSONDecodeError as e:
verification_result = {
'error': f'메타데이터 파싱 오류: {str(e)}',
'confidence': float(confidence),
'threshold': 0.3,
'confidence_level': confidence_info['level'],
'interpretation': confidence_info['message']
}
except Exception as e:
verification_result = {
'error': f'검증 중 오류: {str(e)}',
'confidence': float(confidence),
'threshold': 0.3,
'confidence_level': confidence_info['level'],
'interpretation': confidence_info['message']
}
# 결과 메시지 (한글)
result_msg = f"""
🔍 워터마크 추출 완료!
📊 신뢰도: {confidence:.3f}
{confidence_info['color']} 신뢰도 수준: {confidence_info['level']}
{confidence_info['emoji']} {confidence_info['message']}
💡 {confidence_info['detail']}
"""
if verification_result and 'error' not in verification_result:
status = "✅ 유효" if verification_result['is_valid'] else "❌ 무효"
result_msg += f"""
🛡️ 검증 결과: {status}
📏 임계값: {verification_result['threshold']}
"""
# 메타데이터가 있는 경우 추가 정보
if 'original_metadata' in verification_result:
orig_meta = verification_result['original_metadata']
result_msg += f"""
🆔 원본 사용자: {orig_meta.get('user_id', 'N/A')}
⏰ 생성 시간: {orig_meta.get('timestamp', 'N/A')}
💪 원본 강도: {orig_meta.get('watermark_strength', 'N/A')}
"""
elif verification_result and 'error' in verification_result:
result_msg += f"""
⚠️ 검증 오류: {verification_result['error']}
"""
# 추가 해석 가이드 (한글)
result_msg += f"""
📖 해석 가이드:
• 0.8 이상: 워터마크 확실히 존재 ✅
• 0.6~0.8: 워터마크 존재 가능성 높음 ⚠️
• 0.3~0.6: 워터마크 존재 불확실 ❓
• 0.3 미만: 워터마크 없음 ❌
"""
plt.close()
# JSON 직렬화 검증
verification_json = None
if verification_result:
try:
verification_json = json.dumps(verification_result, indent=2, ensure_ascii=False)
except Exception as e:
verification_json = json.dumps({
'error': f'JSON 직렬화 오류: {str(e)}',
'confidence': float(confidence),
'confidence_level': confidence_info['level'],
'interpretation': confidence_info['message']
}, indent=2, ensure_ascii=False)
return watermark_viz, result_msg, verification_json
except Exception as e:
plt.close()
return None, f"오류 발생: {str(e)}", json.dumps({
'error': f'추출 중 오류: {str(e)}'
}, indent=2, ensure_ascii=False)
def compare_images(original, watermarked):
"""원본과 워터마크된 이미지 비교"""
if original is None or watermarked is None:
return None, "두 이미지가 모두 필요합니다."
try:
# 이미지 크기 맞추기
h1, w1 = original.shape[:2]
h2, w2 = watermarked.shape[:2]
if (h1, w1) != (h2, w2):
watermarked = cv2.resize(watermarked, (w1, h1))
# PSNR 계산
mse = np.mean((original.astype(float) - watermarked.astype(float)) ** 2)
if mse == 0:
psnr = float('inf')
else:
psnr = 20 * np.log10(255.0 / np.sqrt(mse))
# 차이 이미지 생성
diff = np.abs(original.astype(float) - watermarked.astype(float))
diff = (diff / diff.max() * 255).astype(np.uint8)
# 시각화 (제목만 영어로, 폰트 설정)
plt.rcParams['font.family'] = 'DejaVu Sans'
plt.rcParams['font.size'] = 10
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
axes[0].imshow(original)
axes[0].set_title('Original Image', fontweight='bold')
axes[0].axis('off')
axes[1].imshow(watermarked)
axes[1].set_title('Watermarked Image', fontweight='bold')
axes[1].axis('off')
axes[2].imshow(diff, cmap='hot')
axes[2].set_title(f'Difference (PSNR: {psnr:.2f}dB)', fontweight='bold')
axes[2].axis('off')
plt.tight_layout()
result_msg = f"""
📊 이미지 품질 분석:
- PSNR: {psnr:.2f} dB
- MSE: {mse:.2f}
- 이미지 크기: {original.shape}
"""
comparison_fig = plt.gcf()
plt.close()
return comparison_fig, result_msg
except Exception as e:
return None, f"오류 발생: {str(e)}"
# ===== Gradio 인터페이스 구성 =====
def create_gradio_interface():
"""Gradio 인터페이스 생성"""
with gr.Blocks(title="모바일 워터마킹 실험 시스템", theme=gr.themes.Soft()) as demo:
gr.Markdown("""
# 📱 모바일 환경 CNN 기반 워터마킹 시스템
이 시스템은 모바일 환경에 최적화된 실시간 이미지 워터마킹 기술을 실험할 수 있습니다.
## 🔬 주요 기능
- **실시간 워터마크 삽입**: 경량화된 CNN 모델로 빠른 처리
- **적응형 해상도 지원**: 다양한 이미지 크기에 자동 적응
- **워터마크 검증**: 삽입된 워터마크의 추출 및 검증
- **품질 분석**: 원본 대비 화질 변화 측정
""")
with gr.Tabs():
# 탭 1: 워터마크 삽입
with gr.Tab("🔒 워터마크 삽입"):
with gr.Row():
with gr.Column():
embed_input_image = gr.Image(
label="📷 원본 이미지 업로드",
type="numpy"
)
embed_user_id = gr.Textbox(
label="🆔 사용자 ID",
placeholder="예: user123",
value="demo_user"
)
output_format = gr.Radio(
label="📄 출력 포맷",
choices=["PNG", "JPG"],
value="PNG"
)
embed_btn = gr.Button("🔒 워터마크 삽입", variant="primary")
with gr.Column():
embed_output_image = gr.Image(label="🔐 워터마크된 이미지")
embed_result_text = gr.Textbox(
label="📊 처리 결과",
lines=8,
interactive=False
)
download_btn = gr.DownloadButton(
label="💾 워터마크된 이미지 다운로드",
variant="secondary"
)
embed_metadata = gr.JSON(label="📋 메타데이터", visible=False)
embed_btn.click(
fn=embed_watermark_interface,
inputs=[embed_input_image, embed_user_id, output_format],
outputs=[embed_output_image, embed_result_text, embed_metadata, download_btn]
)
# 탭 2: 워터마크 추출 및 검증
with gr.Tab("🔍 워터마크 검증"):
with gr.Row():
with gr.Column():
extract_input_image = gr.Image(
label="🔐 워터마크된 이미지 업로드",
type="numpy"
)
extract_metadata = gr.Textbox(
label="📋 원본 메타데이터 (선택사항)",
placeholder="워터마크 삽입 시 생성된 메타데이터를 붙여넣으세요",
lines=5
)
extract_btn = gr.Button("🔍 워터마크 추출", variant="primary")
with gr.Column():
extract_output_viz = gr.Plot(label="🎨 추출된 워터마크 패턴")
extract_result_text = gr.Textbox(
label="📊 추출 결과",
lines=5,
interactive=False
)
extract_verification = gr.JSON(label="🛡️ 검증 결과", visible=True)
extract_btn.click(
fn=extract_watermark_interface,
inputs=[extract_input_image, extract_metadata],
outputs=[extract_output_viz, extract_result_text, extract_verification]
)
# 탭 3: 이미지 품질 비교
with gr.Tab("📊 품질 분석"):
with gr.Row():
with gr.Column():
compare_original = gr.Image(
label="📷 원본 이미지",
type="numpy"
)
compare_watermarked = gr.Image(
label="🔐 워터마크된 이미지",
type="numpy"
)
compare_btn = gr.Button("📊 품질 비교", variant="primary")
with gr.Column():
compare_output_plot = gr.Plot(label="🔬 비교 분석 결과")
compare_result_text = gr.Textbox(
label="📈 분석 결과",
lines=6,
interactive=False
)
compare_btn.click(
fn=compare_images,
inputs=[compare_original, compare_watermarked],
outputs=[compare_output_plot, compare_result_text]
)
# 탭 4: 시스템 정보
with gr.Tab("ℹ️ 시스템 정보"):
gr.Markdown(f"""
## 🔧 시스템 사양
- **디바이스**: {watermarking_system.device}
- **CNN 아키텍처**: MobileNet 기반 경량화 모델
- **워터마크 크기**: 32x32 픽셀
- **지원 포맷**: JPG, PNG, BMP
## 📈 성능 특징
- **처리 속도**: < 1초 (목표)
- **메모리 효율성**: 모바일 최적화
- **해상도 적응**: 동적 크기 조절
- **견고성**: 압축/변환 공격 저항
## 🎯 사용 방법
1. **워터마크 삽입**: 원본 이미지와 사용자 ID 입력
2. **워터마크 검증**: 의심되는 이미지 업로드 후 추출
3. **품질 분석**: 원본과 워터마크된 이미지 비교
## ⚠️ 주의사항
- 이는 실험용 프로토타입입니다
- 실제 상용 환경에서는 추가 최적화가 필요합니다
- 보안 강화를 위해 더 복잡한 암호화 기법 적용 권장
""")
# 연결 기능: 워터마크 삽입 결과를 검증 탭으로 전달
embed_output_image.change(
fn=lambda x: x,
inputs=[embed_output_image],
outputs=[extract_input_image]
)
embed_metadata.change(
fn=lambda x: json.dumps(x, indent=2) if x else "",
inputs=[embed_metadata],
outputs=[extract_metadata]
)
# 연결 기능: 비교 분석을 위한 이미지 전달
embed_input_image.change(
fn=lambda x: x,
inputs=[embed_input_image],
outputs=[compare_original]
)
embed_output_image.change(
fn=lambda x: x,
inputs=[embed_output_image],
outputs=[compare_watermarked]
)
return demo
# ===== 메인 실행 =====
if __name__ == "__main__":
# Gradio 인터페이스 생성 및 실행
demo = create_gradio_interface()
# Colab 환경에서 실행
demo.launch(
share=True, # 공개 링크 생성
debug=True, # 디버그 모드
server_name="0.0.0.0", # 모든 IP에서 접근 가능
server_port=7860 # 포트 지정
)