ganeshkumar383/AI-Based-Image-Deblurring-App
0
1"""
2CNN Deblurring Module - Deep Learning Based Image Enhancement
3============================================================
4
5CNN inference system for image deblurring with TensorFlow/Keras.
6Includes model architecture, training utilities, and inference pipeline.
7"""
8
9import cv2
10import numpy as np
11import os
12import logging
13from typing import Optional, Tuple, List
14import pickle
15
16# Configure TensorFlow to reduce verbosity
17os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # Reduce TensorFlow logging
18os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' # Disable oneDNN messages
19
20import tensorflow as tf
21from tensorflow import keras
22from tensorflow.keras import layers, Model
23
24# Configure TensorFlow settings
25tf.get_logger().setLevel('ERROR') # Only show errors
26tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
27
28# Configure logging
29logging.basicConfig(level=logging.INFO)
30logger = logging.getLogger(__name__)
31
32class CNNDeblurModel:
33 """CNN-based deblurring model with encoder-decoder architecture"""
34
35 def __init__(self, input_shape: Tuple[int, int, int] = (256, 256, 3)):
36 self.input_shape = input_shape
37 self.model = None
38 self.is_trained = False
39 self.training_history = None
40 self.model_path = "models/cnn_deblur_model.h5"
41 self.dataset_path = "data/training_dataset"
42
43 def build_model(self) -> Model:
44 """
45 Build CNN deblurring model with U-Net like architecture
46
47 Returns:
48 keras.Model: Compiled CNN model
49 """
50 try:
51 # Input layer
52 inputs = keras.Input(shape=self.input_shape)
53
54 # Encoder (Downsampling)
55 conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(inputs)
56 conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv1)
57 pool1 = layers.MaxPooling2D(pool_size=(2, 2))(conv1)
58
59 conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(pool1)
60 conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv2)
61 pool2 = layers.MaxPooling2D(pool_size=(2, 2))(conv2)
62
63 conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(pool2)
64 conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(conv3)
65 pool3 = layers.MaxPooling2D(pool_size=(2, 2))(conv3)
66
67 # Bottleneck
68 conv4 = layers.Conv2D(512, 3, activation='relu', padding='same')(pool3)
69 conv4 = layers.Conv2D(512, 3, activation='relu', padding='same')(conv4)
70
71 # Decoder (Upsampling)
72 up5 = layers.UpSampling2D(size=(2, 2))(conv4)
73 up5 = layers.Conv2D(256, 2, activation='relu', padding='same')(up5)
74 merge5 = layers.concatenate([conv3, up5], axis=3)
75 conv5 = layers.Conv2D(256, 3, activation='relu', padding='same')(merge5)
76 conv5 = layers.Conv2D(256, 3, activation='relu', padding='same')(conv5)
77
78 up6 = layers.UpSampling2D(size=(2, 2))(conv5)
79 up6 = layers.Conv2D(128, 2, activation='relu', padding='same')(up6)
80 merge6 = layers.concatenate([conv2, up6], axis=3)
81 conv6 = layers.Conv2D(128, 3, activation='relu', padding='same')(merge6)
82 conv6 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv6)
83
84 up7 = layers.UpSampling2D(size=(2, 2))(conv6)
85 up7 = layers.Conv2D(64, 2, activation='relu', padding='same')(up7)
86 merge7 = layers.concatenate([conv1, up7], axis=3)
87 conv7 = layers.Conv2D(64, 3, activation='relu', padding='same')(merge7)
88 conv7 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv7)
89
90 # Output layer
91 outputs = layers.Conv2D(3, 1, activation='sigmoid')(conv7)
92
93 # Create model
94 model = Model(inputs=inputs, outputs=outputs)
95
96 # Compile model
97 model.compile(
98 optimizer='adam',
99 loss='mse',
100 metrics=['mae', 'mse']
101 )
102
103 self.model = model
104 logger.info("CNN model built successfully")
105 return model
106
107 except Exception as e:
108 logger.error(f"Error building CNN model: {e}")
109 return None
110
111 def load_model(self, model_path: str) -> bool:
112 """
113 Load pre-trained model from file
114
115 Args:
116 model_path: Path to saved model
117
118 Returns:
119 bool: Success status
120 """
121 try:
122 if os.path.exists(model_path):
123 self.model = keras.models.load_model(model_path)
124 self.is_trained = True
125 logger.info(f"Model loaded from {model_path}")
126 return True
127 else:
128 logger.warning(f"Model file not found: {model_path}")
129 # Build new model as fallback
130 self.build_model()
131 return False
132
133 except Exception as e:
134 logger.error(f"Error loading model: {e}")
135 self.build_model() # Fallback to new model
136 return False
137
138 def save_model(self, model_path: str) -> bool:
139 """
140 Save current model to file
141
142 Args:
143 model_path: Path to save model
144
145 Returns:
146 bool: Success status
147 """
148 try:
149 if self.model is not None:
150 self.model.save(model_path)
151 logger.info(f"Model saved to {model_path}")
152 return True
153 else:
154 logger.error("No model to save")
155 return False
156
157 except Exception as e:
158 logger.error(f"Error saving model: {e}")
159 return False
160
161 def preprocess_image(self, image: np.ndarray) -> np.ndarray:
162 """
163 Preprocess image for CNN input with color preservation
164
165 Args:
166 image: Input image (BGR format)
167
168 Returns:
169 np.ndarray: Preprocessed image
170 """
171 try:
172 # Convert BGR to RGB (preserve original precision)
173 if len(image.shape) == 3 and image.shape[2] == 3:
174 rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
175 else:
176 rgb_image = image
177
178 # Resize to model input size with high-quality interpolation
179 resized = cv2.resize(rgb_image,
180 (self.input_shape[1], self.input_shape[0]),
181 interpolation=cv2.INTER_CUBIC) # Better color preservation
182
183 # Normalize to [0, 1] with high precision
184 normalized = resized.astype(np.float64) / 255.0 # Use float64 for precision
185
186 # Add batch dimension
187 batched = np.expand_dims(normalized, axis=0)
188
189 return batched.astype(np.float32) # Convert to float32 for model
190
191 except Exception as e:
192 logger.error(f"Error preprocessing image: {e}")
193 return np.array([])
194
195 def postprocess_image(self, output: np.ndarray, original_shape: Tuple[int, int]) -> np.ndarray:
196 """
197 Postprocess CNN output to original image format with color preservation
198
199 Args:
200 output: CNN model output
201 original_shape: Original image shape (height, width)
202
203 Returns:
204 np.ndarray: Postprocessed image in BGR format
205 """
206 try:
207 # Remove batch dimension
208 if len(output.shape) == 4:
209 output = output[0]
210
211 # Denormalize from [0, 1] to [0, 255] with high precision
212 denormalized = np.clip(output * 255.0, 0, 255) # Clip before conversion
213 denormalized = np.round(denormalized).astype(np.uint8) # Round to preserve colors
214
215 # Resize to original size with high-quality interpolation
216 resized = cv2.resize(denormalized,
217 (original_shape[1], original_shape[0]),
218 interpolation=cv2.INTER_CUBIC) # Better color preservation
219
220 # Convert RGB back to BGR
221 bgr_image = cv2.cvtColor(resized, cv2.COLOR_RGB2BGR)
222
223 return bgr_image
224
225 except Exception as e:
226 logger.error(f"Error postprocessing image: {e}")
227 return np.zeros((*original_shape, 3), dtype=np.uint8)
228
229 def enhance_image(self, image: np.ndarray) -> np.ndarray:
230 """
231 Enhance image using CNN model
232
233 Args:
234 image: Input blurry image (BGR format)
235
236 Returns:
237 np.ndarray: Enhanced image (BGR format)
238 """
239 try:
240 if self.model is None:
241 logger.warning("No model available, building new model")
242 self.build_model()
243
244 # Store original shape
245 original_shape = image.shape[:2]
246
247 # Preprocess
248 preprocessed = self.preprocess_image(image)
249
250 if preprocessed.size == 0:
251 logger.error("Failed to preprocess image")
252 return image
253
254 # If model is not trained, return enhanced version using traditional methods
255 if not self.is_trained:
256 logger.info("Using fallback enhancement (model not trained)")
257 return self._fallback_enhancement(image)
258
259 # CNN inference
260 enhanced = self.model.predict(preprocessed, verbose=0)
261
262 # Postprocess
263 result = self.postprocess_image(enhanced, original_shape)
264
265 logger.info("CNN enhancement completed")
266 return result
267
268 except Exception as e:
269 logger.error(f"Error in CNN enhancement: {e}")
270 return self._fallback_enhancement(image)
271
272 def _fallback_enhancement(self, image: np.ndarray) -> np.ndarray:
273 """
274 Fallback enhancement when CNN model is not available - preserves original colors
275
276 Args:
277 image: Input image
278
279 Returns:
280 np.ndarray: Enhanced image using color-preserving traditional methods
281 """
282 try:
283 # Method 1: Gentle unsharp masking with color preservation
284 # Create a subtle blur for unsharp masking
285 gaussian = cv2.GaussianBlur(image, (5, 5), 1.0)
286
287 # Apply very gentle unsharp masking to avoid color shifts
288 enhanced = cv2.addWeighted(image, 1.2, gaussian, -0.2, 0)
289
290 # Method 2: Enhance sharpness without changing colors
291 # Convert to float for precision
292 img_float = image.astype(np.float64)
293
294 # Apply high-pass filter for sharpening
295 kernel_sharpen = np.array([[-0.1, -0.1, -0.1],
296 [-0.1, 1.8, -0.1],
297 [-0.1, -0.1, -0.1]])
298
299 # Apply sharpening kernel to each channel separately
300 sharpened_channels = []
301 for i in range(3): # Process each color channel
302 channel = img_float[:, :, i]
303 sharpened_channel = cv2.filter2D(channel, -1, kernel_sharpen)
304 sharpened_channels.append(sharpened_channel)
305
306 sharpened = np.stack(sharpened_channels, axis=2)
307
308 # Combine original with sharpened (gentle blend)
309 result = 0.7 * img_float + 0.3 * sharpened
310
311 # Carefully clip and convert back
312 result = np.clip(result, 0, 255).astype(np.uint8)
313
314 logger.info("Color-preserving fallback enhancement applied")
315 return result
316
317 except Exception as e:
318 logger.error(f"Error in fallback enhancement: {e}")
319 return image
320
321class CNNTrainer:
322 """Training utilities for CNN deblurring model"""
323
324 def __init__(self, model: CNNDeblurModel):
325 self.model = model
326
327 def create_synthetic_data(self, clean_images: List[np.ndarray],
328 blur_types: List[str] = None) -> Tuple[np.ndarray, np.ndarray]:
329 """
330 Create synthetic training data by applying blur to clean images
331
332 Args:
333 clean_images: List of clean images
334 blur_types: Types of blur to apply
335
336 Returns:
337 tuple: (blurred_images, clean_images) for training
338 """
339 if blur_types is None:
340 blur_types = ['gaussian', 'motion', 'defocus']
341
342 blurred_batch = []
343 clean_batch = []
344
345 try:
346 for clean_img in clean_images:
347 # Random blur type
348 blur_type = np.random.choice(blur_types)
349
350 if blur_type == 'gaussian':
351 # Gaussian blur
352 kernel_size = np.random.randint(5, 15)
353 if kernel_size % 2 == 0:
354 kernel_size += 1
355 blurred = cv2.GaussianBlur(clean_img, (kernel_size, kernel_size), 0)
356
357 elif blur_type == 'motion':
358 # Motion blur
359 length = np.random.randint(5, 20)
360 angle = np.random.randint(0, 180)
361 kernel = self._create_motion_kernel(length, angle)
362 blurred = cv2.filter2D(clean_img, -1, kernel)
363
364 else: # defocus
365 # Defocus blur (approximated with Gaussian)
366 sigma = np.random.uniform(1, 5)
367 blurred = cv2.GaussianBlur(clean_img, (0, 0), sigma)
368
369 blurred_batch.append(blurred)
370 clean_batch.append(clean_img)
371
372 return np.array(blurred_batch), np.array(clean_batch)
373
374 except Exception as e:
375 logger.error(f"Error creating synthetic data: {e}")
376 return np.array([]), np.array([])
377
378 def _create_motion_kernel(self, length: int, angle: float) -> np.ndarray:
379 """Create motion blur kernel"""
380 kernel = np.zeros((length, length))
381 center = length // 2
382
383 cos_val = np.cos(np.radians(angle))
384 sin_val = np.sin(np.radians(angle))
385
386 for i in range(length):
387 offset = i - center
388 y = int(center + offset * sin_val)
389 x = int(center + offset * cos_val)
390 if 0 <= y < length and 0 <= x < length:
391 kernel[y, x] = 1
392
393 return kernel / kernel.sum()
394
395 def _load_user_images(self) -> List[np.ndarray]:
396 """Load user's training images from training_dataset folder"""
397 user_images = []
398
399 try:
400 if not os.path.exists(self.dataset_path):
401 return user_images
402
403 # Supported image extensions
404 valid_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
405
406 for filename in os.listdir(self.dataset_path):
407 if any(filename.lower().endswith(ext) for ext in valid_extensions):
408 image_path = os.path.join(self.dataset_path, filename)
409 try:
410 # Load image
411 image = cv2.imread(image_path)
412 if image is not None:
413 # Resize to model input size
414 resized = cv2.resize(image, (self.input_shape[1], self.input_shape[0]))
415 user_images.append(resized)
416 logger.info(f"Loaded user image: {filename}")
417 except Exception as e:
418 logger.warning(f"Failed to load {filename}: {e}")
419
420 logger.info(f"Loaded {len(user_images)} user training images")
421 return user_images
422
423 except Exception as e:
424 logger.error(f"Error loading user images: {e}")
425 return []
426
427 def create_training_dataset(self, num_samples: int = 1000, save_dataset: bool = True) -> Tuple[np.ndarray, np.ndarray]:
428 """
429 Create comprehensive training dataset with various blur types
430 Incorporates user's real training images from data/training_dataset/
431
432 Args:
433 num_samples: Number of training samples to generate
434 save_dataset: Whether to save dataset to disk
435
436 Returns:
437 Tuple[np.ndarray, np.ndarray]: Blurred images and clean targets
438 """
439 try:
440 logger.info(f"Creating training dataset with {num_samples} samples...")
441
442 # Ensure dataset directory exists
443 os.makedirs(self.dataset_path, exist_ok=True)
444
445 # Load user's training images
446 user_images = self._load_user_images()
447
448 all_blurred = []
449 all_clean = []
450
451 # First, process user images if available
452 if user_images:
453 logger.info(f"Processing {len(user_images)} user training images...")
454 for user_img in user_images:
455 # Use user image as clean target multiple times with different blur types
456 for _ in range(3): # Create 3 variations per user image
457 blur_type = np.random.choice(['gaussian', 'motion', 'defocus'])
458
459 if blur_type == 'gaussian':
460 sigma = np.random.uniform(0.5, 3.0)
461 blurred = cv2.GaussianBlur(user_img, (0, 0), sigma)
462 elif blur_type == 'motion':
463 length = np.random.randint(5, 25)
464 angle = np.random.randint(0, 180)
465 kernel = self._create_motion_kernel(length, angle)
466 blurred = cv2.filter2D(user_img, -1, kernel)
467 else: # defocus
468 sigma = np.random.uniform(1.0, 4.0)
469 blurred = cv2.GaussianBlur(user_img, (0, 0), sigma)
470
471 # Add slight noise for realism
472 noise = np.random.normal(0, 3, blurred.shape).astype(np.float32)
473 blurred = np.clip(blurred.astype(np.float32) + noise, 0, 255).astype(np.uint8)
474
475 all_blurred.append(blurred)
476 all_clean.append(user_img)
477
478 # Generate remaining samples with synthetic images
479 remaining_samples = max(0, num_samples - len(all_blurred))
480 if remaining_samples > 0:
481 logger.info(f"Generating {remaining_samples} synthetic training samples...")
482
483 batch_size = 50
484 num_batches = (remaining_samples + batch_size - 1) // batch_size
485
486 for batch_idx in range(num_batches):
487 current_batch_size = min(batch_size, remaining_samples - batch_idx * batch_size)
488
489 # Create synthetic clean images
490 clean_batch = self._generate_clean_images(current_batch_size)
491
492 # Apply various blur types
493 blurred_batch = []
494 for clean_img in clean_batch:
495 blur_type = np.random.choice(['gaussian', 'motion', 'defocus'])
496
497 if blur_type == 'gaussian':
498 sigma = np.random.uniform(0.5, 3.0)
499 blurred = cv2.GaussianBlur(clean_img, (0, 0), sigma)
500 elif blur_type == 'motion':
501 length = np.random.randint(5, 25)
502 angle = np.random.randint(0, 180)
503 kernel = self._create_motion_kernel(length, angle)
504 blurred = cv2.filter2D(clean_img, -1, kernel)
505 else: # defocus
506 sigma = np.random.uniform(1.0, 4.0)
507 blurred = cv2.GaussianBlur(clean_img, (0, 0), sigma)
508
509 # Add slight noise for realism
510 noise = np.random.normal(0, 5, blurred.shape).astype(np.float32)
511 blurred = np.clip(blurred.astype(np.float32) + noise, 0, 255).astype(np.uint8)
512
513 blurred_batch.append(blurred)
514
515 all_blurred.extend(blurred_batch)
516 all_clean.extend(clean_batch)
517
518 if (batch_idx + 1) % 5 == 0:
519 logger.info(f"Generated batch {batch_idx + 1}/{num_batches}")
520
521 # Convert to numpy arrays
522 blurred_dataset = np.array(all_blurred)
523 clean_dataset = np.array(all_clean)
524
525 # Normalize to [0, 1]
526 blurred_dataset = blurred_dataset.astype(np.float32) / 255.0
527 clean_dataset = clean_dataset.astype(np.float32) / 255.0
528
529 logger.info(f"Dataset created: {blurred_dataset.shape} blurred, {clean_dataset.shape} clean")
530
531 # Save dataset if requested
532 if save_dataset:
533 np.save(os.path.join(self.dataset_path, 'blurred_images.npy'), blurred_dataset)
534 np.save(os.path.join(self.dataset_path, 'clean_images.npy'), clean_dataset)
535 logger.info(f"Dataset saved to {self.dataset_path}")
536
537 return blurred_dataset, clean_dataset
538
539 except Exception as e:
540 logger.error(f"Error creating training dataset: {e}")
541 return np.array([]), np.array([])
542
543 def _generate_clean_images(self, num_images: int) -> List[np.ndarray]:
544 """Generate synthetic clean images for training"""
545 clean_images = []
546
547 for _ in range(num_images):
548 # Create random patterns and shapes
549 img = np.zeros((self.input_shape[0], self.input_shape[1], 3), dtype=np.uint8)
550
551 # Random background
552 bg_color = np.random.randint(0, 255, 3)
553 img[:] = bg_color
554
555 # Add random shapes
556 num_shapes = np.random.randint(3, 8)
557 for _ in range(num_shapes):
558 shape_type = np.random.choice(['rectangle', 'circle', 'line'])
559 color = np.random.randint(0, 255, 3).tolist()
560
561 if shape_type == 'rectangle':
562 pt1 = (np.random.randint(0, img.shape[1]//2), np.random.randint(0, img.shape[0]//2))
563 pt2 = (np.random.randint(img.shape[1]//2, img.shape[1]),
564 np.random.randint(img.shape[0]//2, img.shape[0]))
565 cv2.rectangle(img, pt1, pt2, color, -1)
566
567 elif shape_type == 'circle':
568 center = (np.random.randint(0, img.shape[1]), np.random.randint(0, img.shape[0]))
569 radius = np.random.randint(10, 50)
570 cv2.circle(img, center, radius, color, -1)
571
572 else: # line
573 pt1 = (np.random.randint(0, img.shape[1]), np.random.randint(0, img.shape[0]))
574 pt2 = (np.random.randint(0, img.shape[1]), np.random.randint(0, img.shape[0]))
575 thickness = np.random.randint(1, 5)
576 cv2.line(img, pt1, pt2, color, thickness)
577
578 # Add random text
579 if np.random.random() > 0.5:
580 text = ''.join(np.random.choice(list('ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'),
581 np.random.randint(3, 8)))
582 font = cv2.FONT_HERSHEY_SIMPLEX
583 font_scale = np.random.uniform(0.5, 2.0)
584 color = np.random.randint(0, 255, 3).tolist()
585 thickness = np.random.randint(1, 3)
586 position = (np.random.randint(0, img.shape[1]//2), np.random.randint(20, img.shape[0]))
587 cv2.putText(img, text, position, font, font_scale, color, thickness)
588
589 clean_images.append(img)
590
591 return clean_images
592
593 def load_existing_dataset(self) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
594 """Load existing dataset from disk"""
595 try:
596 blurred_path = os.path.join(self.dataset_path, 'blurred_images.npy')
597 clean_path = os.path.join(self.dataset_path, 'clean_images.npy')
598
599 if os.path.exists(blurred_path) and os.path.exists(clean_path):
600 blurred_data = np.load(blurred_path)
601 clean_data = np.load(clean_path)
602 logger.info(f"Loaded existing dataset: {blurred_data.shape} samples")
603 return blurred_data, clean_data
604 else:
605 logger.info("No existing dataset found")
606 return None, None
607
608 except Exception as e:
609 logger.error(f"Error loading existing dataset: {e}")
610 return None, None
611
612 def train_model(self,
613 epochs: int = 20,
614 batch_size: int = 16,
615 validation_split: float = 0.2,
616 use_existing_dataset: bool = True,
617 num_training_samples: int = 1000) -> bool:
618 """
619 Train the CNN model with comprehensive dataset
620
621 Args:
622 epochs: Number of training epochs
623 batch_size: Training batch size
624 validation_split: Fraction of data for validation
625 use_existing_dataset: Whether to use existing saved dataset
626 num_training_samples: Number of samples to generate if creating new dataset
627
628 Returns:
629 bool: Training success status
630 """
631 try:
632 logger.info("Starting CNN model training...")
633
634 # Build model if not exists
635 if self.model is None:
636 self.build_model()
637
638 # Load or create dataset
639 if use_existing_dataset:
640 blurred_data, clean_data = self.load_existing_dataset()
641 if blurred_data is None:
642 logger.info("Creating new dataset...")
643 blurred_data, clean_data = self.create_training_dataset(num_training_samples)
644 else:
645 logger.info("Creating new dataset...")
646 blurred_data, clean_data = self.create_training_dataset(num_training_samples)
647
648 if len(blurred_data) == 0:
649 logger.error("Failed to create/load training dataset")
650 return False
651
652 logger.info(f"Training on {len(blurred_data)} samples")
653
654 # Setup callbacks
655 callbacks = [
656 keras.callbacks.EarlyStopping(
657 monitor='val_loss',
658 patience=5,
659 restore_best_weights=True
660 ),
661 keras.callbacks.ReduceLROnPlateau(
662 monitor='val_loss',
663 factor=0.5,
664 patience=3,
665 min_lr=1e-7
666 ),
667 keras.callbacks.ModelCheckpoint(
668 filepath=self.model_path,
669 monitor='val_loss',
670 save_best_only=True,
671 save_weights_only=False
672 )
673 ]
674
675 # Train model
676 self.training_history = self.model.fit(
677 blurred_data, clean_data,
678 epochs=epochs,
679 batch_size=batch_size,
680 validation_split=validation_split,
681 callbacks=callbacks,
682 verbose=1
683 )
684
685 # Save final model
686 self.save_model(self.model_path)
687 self.is_trained = True
688
689 # Save training history
690 history_path = self.model_path.replace('.h5', '_history.pkl')
691 with open(history_path, 'wb') as f:
692 pickle.dump(self.training_history.history, f)
693
694 logger.info("Training completed successfully!")
695 logger.info(f"Model saved to: {self.model_path}")
696
697 # Print training summary
698 final_loss = self.training_history.history['loss'][-1]
699 final_val_loss = self.training_history.history['val_loss'][-1]
700 logger.info(f"Final training loss: {final_loss:.4f}")
701 logger.info(f"Final validation loss: {final_val_loss:.4f}")
702
703 return True
704
705 except Exception as e:
706 logger.error(f"Error during training: {e}")
707 return False
708
709 def evaluate_model(self, test_images: np.ndarray = None, test_targets: np.ndarray = None) -> dict:
710 """
711 Evaluate model performance on test data
712
713 Args:
714 test_images: Test images (if None, creates synthetic test set)
715 test_targets: Test targets (if None, creates synthetic test set)
716
717 Returns:
718 dict: Evaluation metrics
719 """
720 try:
721 if self.model is None or not self.is_trained:
722 logger.error("Model not trained. Train the model first.")
723 return {}
724
725 # Create test data if not provided
726 if test_images is None or test_targets is None:
727 logger.info("Creating test dataset...")
728 test_images, test_targets = self.create_training_dataset(num_samples=100, save_dataset=False)
729
730 # Evaluate
731 results = self.model.evaluate(test_images, test_targets, verbose=0)
732
733 metrics = {
734 'loss': results[0],
735 'mae': results[1],
736 'mse': results[2]
737 }
738
739 logger.info("Model Evaluation Results:")
740 for metric, value in metrics.items():
741 logger.info(f" {metric}: {value:.4f}")
742
743 return metrics
744
745 except Exception as e:
746 logger.error(f"Error during evaluation: {e}")
747 return {}
748
749# Convenience functions
750def load_cnn_model(model_path: str = "models/cnn_model.h5") -> CNNDeblurModel:
751 """
752 Load CNN deblurring model
753
754 Args:
755 model_path: Path to model file
756
757 Returns:
758 CNNDeblurModel: Loaded model instance
759 """
760 model = CNNDeblurModel()
761 model.load_model(model_path)
762 return model
763
764def enhance_with_cnn(image: np.ndarray, model_path: str = "models/cnn_model.h5") -> np.ndarray:
765 """
766 Enhance image using CNN model
767
768 Args:
769 image: Input image
770 model_path: Path to model file
771
772 Returns:
773 np.ndarray: Enhanced image
774 """
775 model = load_cnn_model(model_path)
776 return model.enhance_image(image)
777
778# Training utility functions
779def train_new_model(num_samples: int = 1000, epochs: int = 20, input_shape: Tuple[int, int, int] = (256, 256, 3)):
780 """
781 Train a new CNN deblurring model from scratch
782
783 Args:
784 num_samples: Number of training samples to generate
785 epochs: Number of training epochs
786 input_shape: Input image shape
787
788 Returns:
789 CNNDeblurModel: Trained model
790 """
791 print("🚀 Training New CNN Deblurring Model")
792 print("=" * 50)
793
794 # Ensure directories exist
795 os.makedirs("models", exist_ok=True)
796 os.makedirs("data/training_dataset", exist_ok=True)
797
798 # Initialize model
799 model = CNNDeblurModel(input_shape=input_shape)
800
801 # Train model
802 success = model.train_model(
803 epochs=epochs,
804 batch_size=16,
805 validation_split=0.2,
806 use_existing_dataset=True,
807 num_training_samples=num_samples
808 )
809
810 if success:
811 print("✅ Training completed successfully!")
812
813 # Evaluate model
814 metrics = model.evaluate_model()
815 if metrics:
816 print(f"📊 Model Performance:")
817 print(f" Loss: {metrics['loss']:.4f}")
818 print(f" MAE: {metrics['mae']:.4f}")
819 print(f" MSE: {metrics['mse']:.4f}")
820
821 return model
822 else:
823 print("❌ Training failed!")
824 return None
825
826def quick_train():
827 """Quick training with default parameters"""
828 return train_new_model(num_samples=500, epochs=10)
829
830def full_train():
831 """Full training with comprehensive dataset"""
832 return train_new_model(num_samples=2000, epochs=30)
833
834# Example usage and testing
835if __name__ == "__main__":
836 import argparse
837
838 parser = argparse.ArgumentParser(description='CNN Deblurring Module')
839 parser.add_argument('--train', action='store_true', help='Train the model')
840 parser.add_argument('--quick-train', action='store_true', help='Quick training (500 samples, 10 epochs)')
841 parser.add_argument('--full-train', action='store_true', help='Full training (2000 samples, 30 epochs)')
842 parser.add_argument('--samples', type=int, default=1000, help='Number of training samples')
843 parser.add_argument('--epochs', type=int, default=20, help='Number of training epochs')
844 parser.add_argument('--test', action='store_true', help='Test the model')
845
846 args = parser.parse_args()
847
848 print("🎯 CNN Deblurring Module")
849 print("=" * 30)
850
851 if args.quick_train:
852 print("🚀 Quick Training Mode")
853 model = quick_train()
854
855 elif args.full_train:
856 print("🚀 Full Training Mode")
857 model = full_train()
858
859 elif args.train:
860 print(f"🚀 Custom Training Mode")
861 model = train_new_model(num_samples=args.samples, epochs=args.epochs)
862
863 elif args.test:
864 print("🧪 Testing Mode")
865
866 # Create test image
867 test_image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
868
869 # Initialize model
870 cnn_model = CNNDeblurModel()
871
872 # Try to load existing model
873 if cnn_model.load_model(cnn_model.model_path):
874 print(f"✅ Loaded existing trained model")
875 else:
876 print(f"ℹ️ No trained model found, building new model")
877 cnn_model.build_model()
878
879 print(f"Model input shape: {cnn_model.input_shape}")
880 print(f"Model built: {cnn_model.model is not None}")
881 print(f"Model trained: {cnn_model.is_trained}")
882
883 # Test enhancement
884 enhanced = cnn_model.enhance_image(test_image)
885 print(f"Original shape: {test_image.shape}")
886 print(f"Enhanced shape: {enhanced.shape}")
887
888 if cnn_model.is_trained:
889 # Evaluate on test data
890 metrics = cnn_model.evaluate_model()
891 if metrics:
892 print("📊 Model Performance:")
893 for metric, value in metrics.items():
894 print(f" {metric}: {value:.4f}")
895
896 else:
897 print("ℹ️ Usage options:")
898 print(" --test Test existing model or build new one")
899 print(" --quick-train Quick training (500 samples, 10 epochs)")
900 print(" --full-train Full training (2000 samples, 30 epochs)")
901 print(" --train Custom training (use --samples and --epochs)")
902 print("\nExamples:")
903 print(" python -m modules.cnn_deblurring --test")
904 print(" python -m modules.cnn_deblurring --quick-train")
905 print(" python -m modules.cnn_deblurring --train --samples 1500 --epochs 25")
906
907 print("\n🎯 CNN deblurring module ready!")