Upload train_model.py

train_model.py

@@ -0,0 +1,73 @@
+#train.py
+
+import os
+import tensorflow as tf
+import matplotlib.pyplot as plt
+
+# Import the new data loader and the existing model builder
+from utils import load_div2k_data
+from model import build_enhanced_model, psnr
+
+# --- 1. Training Configuration ---
+BATCH_SIZE = 16 # Smaller batch size for larger images to fit in GPU memory
+EPOCHS = 30     # Fewer epochs, as each one takes longer. Increase for higher quality.
+
+# --- 2. Load the Dataset ---
+train_ds, valid_ds, ds_info = load_div2k_data(batch_size=BATCH_SIZE)
+
+# Calculate steps per epoch
+steps_per_epoch = ds_info.splits['train'].num_examples // BATCH_SIZE
+validation_steps = ds_info.splits['validation'].num_examples // BATCH_SIZE
+
+# --- 3. Build the Model for 128x128 Input ---
+INPUT_SHAPE = (128, 128, 3)
+model = build_enhanced_model(input_shape=INPUT_SHAPE)
+model.summary()
+
+# --- 4. Train the Model ---
+print("\nStarting model training on 128x128 images...")
+history = model.fit(
+    train_ds,
+    epochs=EPOCHS,
+    steps_per_epoch=steps_per_epoch,
+    validation_data=valid_ds,
+    validation_steps=validation_steps
+)
+print("Training finished.")
+
+# --- 5. Save the New Model ---
+if not os.path.exists('models'):
+    os.makedirs('models')
+
+model_path = 'models/sr_128_model.h5'
+model.save(model_path)
+print(f"✅ Model for 128x128 saved to {model_path}")
+
+# --- 6. Visualize a Test Result ---
+print("\nVisualizing a sample prediction...")
+# Get one batch from the validation dataset to visualize
+for lr_batch, hr_batch in valid_ds.take(1):
+    # Take the first image from the batch
+    lr_image = lr_batch[0]
+    hr_image = hr_batch[0]
+
+    # Predict
+    pred_image = model.predict(tf.expand_dims(lr_image, axis=0))[0]
+
+    # Plot
+    plt.figure(figsize=(15, 6))
+    plt.subplot(1, 3, 1)
+    plt.imshow(lr_image)
+    plt.title('Low-Res Input (128x128 Upscaled)')
+    plt.axis('off')
+
+    plt.subplot(1, 3, 2)
+    plt.imshow(tf.clip_by_value(pred_image, 0, 1)) # Clip values to [0,1] for display
+    plt.title('AI Super-Resolved Output')
+    plt.axis('off')
+
+    plt.subplot(1, 3, 3)
+    plt.imshow(hr_image)
+    plt.title('Original High-Resolution')
+    plt.axis('off')
+    plt.show()