ImageWatermarking/mainwindow.py at main · wmapro1001/ImageWatermarking · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
from PyQt5.QtWidgets import QWidget, QLabel, QLineEdit, QPushButton,  QVBoxLayout, QHBoxLayout
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIcon, QFont, QFontDatabase, QPainter,QImage,QPixmap
from PIL import Image
import cv2
import numpy as np

class MainWindow(QWidget):
	def __init__(self):
		super().__init__()

		self.setWindowIcon(QIcon('./assets/logo.png'))
		self.setWindowTitle("Invisible Image Watermarking")
		self.setFixedHeight(600)
		self.setObjectName("main_window")
		self.lblOriImg = QLabel()
		self.lblOriImg.setFixedWidth(500)
		self.lblOriImg.setStyleSheet("border: 1px solid black")
		self.lblWatermark = QLabel()
		self.lblWatermark.setAlignment(Qt.AlignCenter)
		self.lblWatermarkEx = QLabel()
		self.lblWatermarkEx.setAlignment(Qt.AlignCenter)
		self.lblResImg = QLabel()
		self.lblResImg.setFixedWidth(500)
		self.lblResImg.setStyleSheet("border: 1px solid black")

		self.btnLoadWM = QPushButton("Load Watermark")
		self.btnLoadWM.setFixedWidth(200)
		self.btnLoadImg = QPushButton("Load Original Image")
		self.btnEmbedWM = QPushButton("Embed Watermark")
		self.btnExtractWM = QPushButton("Extract Watermark")
		self.btnEmbedWM.setFixedWidth(200)

		self.mainLayout = QHBoxLayout()
		self.mainLayout.addWidget(self.lblOriImg)
		self.btnLayout = QVBoxLayout()
		self.btnLayout.addStretch(1)
		self.btnLayout.addWidget(self.lblWatermark)
		self.btnLayout.addWidget(self.btnLoadImg)
		self.btnLayout.addWidget(self.btnLoadWM)
		self.btnLayout.addWidget(self.btnEmbedWM)
		self.btnLayout.addWidget(self.btnExtractWM)
		self.btnLayout.addWidget(self.lblWatermarkEx)
		self.btnLayout.addStretch(1)
		self.mainLayout.addLayout(self.btnLayout)
		self.mainLayout.addWidget(self.lblResImg)

		self.setLayout(self.mainLayout)

		self.btnLoadWM.clicked.connect(self.loadWM)
		self.btnEmbedWM.clicked.connect(self.embedWM)
		self.btnLoadImg.clicked.connect(self.loadOriImg)
		self.btnExtractWM.clicked.connect(self.extractWM)

		self.oriImg = None
		self.wmImg = None
		self.watermarked_img = None

	def loadOriImg(self):
		print("clicked load img btn")
		file_name, _ = QFileDialog.getOpenFileName(self, "Open Original Image", "", "Image Files (*.png *.jpg *.bmp *.jpeg)")
		if file_name:
			pixmap = QPixmap(file_name)
			self.lblOriImg.setPixmap(pixmap)
			self.oriImg = cv2.imread(file_name)

	def loadWM(self):
		print("clicked load watermark")
		file_name, _ = QFileDialog.getOpenFileName(self, "Open Watermark Image", "", "Image Files (*.png *.jpg *.bmp *.jpeg)")
		if file_name:
			pixmap = QPixmap(file_name)
			self.lblWatermark.setPixmap(pixmap)
			self.wmImg = cv2.imread(file_name)


	def embedWM(self):
		print("clicked embed watermark")
		self.watermarked_img = self.embed_lsb_watermark(self.oriImg, self.wmImg, channel='R')
		cv2.imwrite("watermarked.png", self.watermarked_img)
		pixmap = self.cv2_to_pixmap(self.watermarked_img)
		self.lblResImg.setPixmap(pixmap)

	#embed watermark to top-left 64*64 part
	def embed_lsb_watermark(self, host_img, watermark_img, channel='R'):

		# Ensure watermark is 64x64 and binary (0/1)
	    watermark_img = cv2.resize(watermark_img, (64, 64))
	    watermark_gray = cv2.cvtColor(watermark_img, cv2.COLOR_BGR2GRAY)
	    _, watermark_binary = cv2.threshold(watermark_gray, 127, 1, cv2.THRESH_BINARY)

	    # Select channel index
	    channel_index = {'B': 0, 'G': 1, 'R': 2}[channel]

	    # Make a copy to not modify original
	    watermarked_img = host_img.copy()

	    # Embed watermark in top-left 64x64 area
	    for y in range(64):
	        for x in range(64):
	            pixel_value = watermarked_img[y, x, channel_index]
	            watermarked_img[y, x, channel_index] = (pixel_value & 0b11111110) | watermark_binary[y, x]

	    return watermarked_img

	def cv2_to_pixmap(self,cv_img):
	    # Ensure image is in BGR format (default in OpenCV)
	    if len(cv_img.shape) == 3:  # Color image
	        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
	        h, w, ch = cv_img.shape
	        bytes_per_line = ch * w
	        qimage = QImage(cv_img.data, w, h, bytes_per_line, QImage.Format_RGB888)
	    else:         # Grayscale image
	        h, w = cv_img.shape
	        bytes_per_line = w
	        qimage = QImage(cv_img.data, w, h, bytes_per_line, QImage.Format_Grayscale8)
	    return QPixmap.fromImage(qimage)


	def extractWM(self):
		print("extract WM clicked")
		watermark_extracted = self.extract_lsb_watermark(self.watermarked_img, size=64, channel='R')
		# cv2.imshow("Extracted Watermark", watermark_extracted)
		pixmap = self.cv2_to_pixmap(watermark_extracted)
		self.lblWatermarkEx.setPixmap(pixmap)

	def extract_lsb_watermark(self, watermarked_img, size=64 , channel='R'):
		# Map channel letter to index
	    channel_index = {'B': 0, 'G': 1, 'R': 2}[channel]

	    # Extract top-left area
	    lsb_bits = watermarked_img[0:size, 0:size, channel_index] & 1  # get LSB

	    # Convert 0/1 to 0/255 for display
	    watermark_extracted = (lsb_bits * 255).astype(np.uint8)
	    return watermark_extracted