Coverage for evaluation / tools / image_editor.py: 97.01%

1from PIL import Image, ImageFilter, ImageEnhance, ImageOps, ImageDraw 

2import os 

3import argparse 

4import sys 

5import numpy as np 

6import random 

7 

8 

9class ImageEditor: 

10 def __init__(self): 

11 pass 

12 

13 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

14 pass 

15 

16class JPEGCompression(ImageEditor): 

17 def __init__(self, quality: int = 95): 

18 super().__init__() 

19 self.quality = quality 

20 

21 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

22 image.save(f"temp.jpg", quality=self.quality) 

23 compressed_image = Image.open(f"temp.jpg") 

24 os.remove(f"temp.jpg") 

25 return compressed_image 

26 

27class Rotation(ImageEditor): 

28 def __init__(self, angle: int = 30, expand: bool = False): 

29 super().__init__() 

30 self.angle = angle 

31 self.expand = expand 

32 

33 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

34 return image.rotate(self.angle, expand=self.expand) 

35 

36class CrSc(ImageEditor): 

37 def __init__(self, crop_ratio: float = 0.8): 

38 super().__init__() 

39 self.crop_ratio = crop_ratio 

40 

41 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

42 width, height = image.size 

43 new_w = int(width * self.crop_ratio) 

44 new_h = int(height * self.crop_ratio) 

45 

46 left = (width - new_w) // 2 

47 top = (height - new_h) // 2 

48 right = left + new_w 

49 bottom = top + new_h 

50 

51 return image.crop((left, top, right, bottom)).resize((width, height)) 

52 

53class GaussianBlurring(ImageEditor): 

54 def __init__(self, radius: int = 2): 

55 super().__init__() 

56 self.radius = radius 

57 

58 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

59 return image.filter(ImageFilter.GaussianBlur(self.radius)) 

60 

61class GaussianNoise(ImageEditor): 

62 def __init__(self, sigma: float = 25.0): 

63 super().__init__() 

64 self.sigma = sigma 

65 

66 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

67 img = image.convert("RGB") 

68 arr = np.array(img).astype(np.float32) 

69 

70 noise = np.random.normal(0, self.sigma, arr.shape) 

71 noisy_arr = np.clip(arr + noise, 0, 255).astype(np.uint8) 

72 

73 return Image.fromarray(noisy_arr) 

74 

75class Brightness(ImageEditor): 

76 def __init__(self, factor: float = 1.2): 

77 super().__init__() 

78 self.factor = factor 

79 

80 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

81 enhancer = ImageEnhance.Brightness(image) 

82 return enhancer.enhance(self.factor) 

83 

84class Mask(ImageEditor): 

85 def __init__(self, mask_ratio: float = 0.1, num_masks: int = 5): 

86 super().__init__() 

87 self.mask_ratio = mask_ratio 

88 self.num_masks = num_masks 

89 

90 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

91 img = image.copy() 

92 draw = ImageDraw.Draw(img) 

93 width, height = img.size 

94 

95 for _ in range(self.num_masks): 

96 max_mask_width = int(width * self.mask_ratio) 

97 max_mask_height = int(height * self.mask_ratio) 

98 

99 mask_width = random.randint(max_mask_width // 2, max_mask_width) 

100 mask_height = random.randint(max_mask_height // 2, max_mask_height) 

101 

102 x = random.randint(0, width - mask_width) 

103 y = random.randint(0, height - mask_height) 

104 

105 draw.rectangle([x, y, x + mask_width, y + mask_height], fill='black') 

106 

107 return img 

108 

109class Overlay(ImageEditor): 

110 def __init__(self, num_strokes: int = 10, stroke_width: int = 5, stroke_type: str = 'random'): 

111 super().__init__() 

112 self.num_strokes = num_strokes 

113 self.stroke_width = stroke_width 

114 self.stroke_type = stroke_type 

115 

116 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

117 img = image.copy() 

118 draw = ImageDraw.Draw(img) 

119 width, height = img.size 

120 

121 for _ in range(self.num_strokes): 

122 start_x = random.randint(0, width) 

123 start_y = random.randint(0, height) 

124 num_points = random.randint(3, 8) 

125 points = [(start_x, start_y)] 

126 

127 for i in range(num_points - 1): 

128 last_x, last_y = points[-1] 

129 max_step = min(width, height) // 4 

130 new_x = max(0, min(width, last_x + random.randint(-max_step, max_step))) 

131 new_y = max(0, min(height, last_y + random.randint(-max_step, max_step))) 

132 points.append((new_x, new_y)) 

133 

134 if self.stroke_type == 'random': 

135 color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) 

136 elif self.stroke_type == 'black': 

137 color = (0, 0, 0) 

138 elif self.stroke_type == 'white': 

139 color = (255, 255, 255) 

140 else: 

141 color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) 

142 

143 draw.line(points, fill=color, width=self.stroke_width) 

144 

145 return img 

146 

147class AdaptiveNoiseInjection(ImageEditor): 

148 def __init__(self, intensity: float = 0.5, auto_select: bool = True): 

149 super().__init__() 

150 self.intensity = intensity 

151 self.auto_select = auto_select 

152 

153 def _analyze_image_features(self, img_array): 

154 if len(img_array.shape) == 3: 

155 gray = np.mean(img_array, axis=2) 

156 else: 

157 gray = img_array 

158 

159 brightness_mean = np.mean(gray) 

160 brightness_std = np.std(gray) 

161 

162 sobel_x = np.abs(np.diff(gray, axis=1, prepend=gray[:, :1])) 

163 sobel_y = np.abs(np.diff(gray, axis=0, prepend=gray[:1, :])) 

164 edge_density = np.mean(sobel_x + sobel_y) 

165 

166 kernel_size = 5 

167 texture_complexity = 0 

168 h, w = gray.shape 

169 for i in range(0, h - kernel_size, kernel_size): 

170 for j in range(0, w - kernel_size, kernel_size): 

171 patch = gray[i:i+kernel_size, j:j+kernel_size] 

172 texture_complexity += np.std(patch) 

173 texture_complexity /= ((h // kernel_size) * (w // kernel_size)) 

174 

175 return { 

176 'brightness_mean': brightness_mean, 

177 'brightness_std': brightness_std, 

178 'edge_density': edge_density, 

179 'texture_complexity': texture_complexity 

180 } 

181 

182 def _select_noise_type(self, features): 

183 brightness = features['brightness_mean'] 

184 edge_density = features['edge_density'] 

185 texture = features['texture_complexity'] 

186 

187 if brightness < 80: 

188 return 'gaussian' 

189 elif edge_density > 30: 

190 return 'salt_pepper' 

191 elif texture > 20: 

192 return 'speckle' 

193 else: 

194 return 'poisson' 

195 

196 def _add_gaussian_noise(self, img_array, sigma): 

197 noise = np.random.normal(0, sigma, img_array.shape) 

198 noisy = np.clip(img_array + noise, 0, 255) 

199 return noisy.astype(np.uint8) 

200 

201 def _add_salt_pepper_noise(self, img_array, amount): 

202 noisy = img_array.copy() 

203 h, w = img_array.shape[:2] 

204 num_pixels = h * w 

205 

206 num_salt = int(amount * num_pixels * 0.5) 

207 salt_coords_y = np.random.randint(0, h, num_salt) 

208 salt_coords_x = np.random.randint(0, w, num_salt) 

209 noisy[salt_coords_y, salt_coords_x] = 255 

210 

211 num_pepper = int(amount * num_pixels * 0.5) 

212 pepper_coords_y = np.random.randint(0, h, num_pepper) 

213 pepper_coords_x = np.random.randint(0, w, num_pepper) 

214 noisy[pepper_coords_y, pepper_coords_x] = 0 

215 

216 return np.clip(noisy, 0, 255).astype(np.uint8) 

217 

218 def _add_poisson_noise(self, img_array): 

219 vals = len(np.unique(img_array)) 

220 vals = 2 ** np.ceil(np.log2(vals)) 

221 noisy = np.random.poisson(img_array * vals) / float(vals) 

222 return np.clip(noisy, 0, 255).astype(np.uint8) 

223 

224 def _add_speckle_noise(self, img_array, variance): 

225 noise = np.random.randn(*img_array.shape) * variance 

226 noisy = img_array + img_array * noise 

227 return np.clip(noisy, 0, 255).astype(np.uint8) 

228 

229 def edit(self, image: Image.Image, prompt: str = None) -> Image.Image: 

230 img = image.convert("RGB") 

231 img_array = np.array(img).astype(np.float32) 

232 

233 features = self._analyze_image_features(img_array) 

234 

235 if self.auto_select: 

236 noise_type = self._select_noise_type(features) 

237 

238 if noise_type == 'gaussian': 

239 sigma = 40 * self.intensity 

240 noisy_array = self._add_gaussian_noise(img_array, sigma) 

241 elif noise_type == 'salt_pepper': 

242 amount = 0.15 * self.intensity 

243 noisy_array = self._add_salt_pepper_noise(img_array, amount) 

244 elif noise_type == 'poisson': 

245 noisy_array = self._add_poisson_noise(img_array) 

246 blend_factor = min(0.8, self.intensity * 1.5) 

247 noisy_array = np.clip( 

248 img_array * (1 - blend_factor) + noisy_array * blend_factor, 

249 0, 255 

250 ).astype(np.uint8) 

251 else: 

252 variance = 0.5 * self.intensity 

253 noisy_array = self._add_speckle_noise(img_array, variance) 

254 else: 

255 weight = 0.25 

256 noisy_array = img_array.copy() 

257 

258 gaussian = self._add_gaussian_noise(img_array, 30 * self.intensity) 

259 noisy_array = noisy_array * (1 - weight) + gaussian * weight 

260 

261 salt_pepper = self._add_salt_pepper_noise(img_array, 0.08 * self.intensity) 

262 noisy_array = noisy_array * (1 - weight) + salt_pepper * weight 

263 

264 poisson = self._add_poisson_noise(img_array) 

265 noisy_array = noisy_array * (1 - weight) + poisson * weight 

266 

267 speckle = self._add_speckle_noise(img_array, 0.4 * self.intensity) 

268 noisy_array = noisy_array * (1 - weight) + speckle * weight 

269 

270 noisy_array = np.clip(noisy_array, 0, 255).astype(np.uint8) 

271 

272 return Image.fromarray(noisy_array)