sakrit28/Food_Status_Classification

1# train_model.py
2
3import tensorflow as tf
4from tensorflow.keras.preprocessing.image import ImageDataGenerator
5from tensorflow.keras.models import Sequential
6from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
7import matplotlib.pyplot as plt
8import os
9
10# --- 1. Define Constants ---
11# You can change these. Smaller sizes (like 150x150) train faster.
12IMG_WIDTH = 150
13IMG_HEIGHT = 150
14
15# How many images to process at once
16BATCH_SIZE = 32 
17
18# How many times to go over the entire dataset
19EPOCHS = 50  # Start with 20. If accuracy is low, you can increase this later.
20
21base_dir = 'dataset'
22train_dir = os.path.join(base_dir, 'train')
23validation_dir = os.path.join(base_dir, 'validation')
24
25# --- 2. Get Image Counts (for calculating steps) ---
26# This is important for telling Keras how much data to expect in each epoch.
27try:
28    num_train_fresh = len(os.listdir(os.path.join(train_dir, 'fresh')))
29    num_train_rotten = len(os.listdir(os.path.join(train_dir, 'rotten')))
30    num_val_fresh = len(os.listdir(os.path.join(validation_dir, 'fresh')))
31    num_val_rotten = len(os.listdir(os.path.join(validation_dir, 'rotten')))
32except FileNotFoundError:
33    print("Error: Make sure your 'dataset' folder is structured correctly:")
34    print("dataset/train/fresh, dataset/train/rotten, etc.")
35    exit()
36
37total_train = num_train_fresh + num_train_rotten
38total_val = num_val_fresh + num_val_rotten
39
40print(f"Total training images: {total_train}")
41print(f"Total validation images: {total_val}")
42
43
44# --- 3. Data Preprocessing & Augmentation ---
45
46# Create a tool (ImageDataGenerator) to automatically prepare our images.
47# 'rescale' converts pixel values from [0, 255] to [0, 1], which models prefer.
48train_datagen = ImageDataGenerator(
49    rescale=1./255,
50    rotation_range=40,      # Randomly rotate images
51    width_shift_range=0.2,  # Randomly shift images horizontally
52    height_shift_range=0.2, # Randomly shift images vertically
53    shear_range=0.2,        # 'Shear' or 'slant' images
54    zoom_range=0.2,         # Randomly zoom in
55    horizontal_flip=True,   # Randomly flip images horizontally
56    fill_mode='nearest'     # How to fill in new pixels after a rotation/shift
57)
58
59# For the validation data, we ONLY rescale it. 
60# We don't augment it because we want to test the model on 'real' data.
61validation_datagen = ImageDataGenerator(rescale=1./255)
62
63# --- 4. Load Data from Directories ---
64
65# 'flow_from_directory' automatically finds our images, resizes them,
66# and sorts them into batches based on the folder structure.
67train_generator = train_datagen.flow_from_directory(
68    train_dir,
69    target_size=(IMG_WIDTH, IMG_HEIGHT), # Resize all images to this
70    batch_size=BATCH_SIZE,
71    class_mode='binary'  # 'binary' because we have only two classes (fresh/rotten)
72)
73
74validation_generator = validation_datagen.flow_from_directory(
75    validation_dir,
76    target_size=(IMG_WIDTH, IMG_HEIGHT),
77    batch_size=BATCH_SIZE,
78    class_mode='binary'
79)
80
81# CRITICAL: Note which class is 0 and which is 1
82print("Class Indices (this is important!):", train_generator.class_indices)
83
84
85# --- 5. Build the CNN Model ---
86# We build the model layer by layer.
87model = Sequential([
88    # Layer 1: Convolution + Pooling
89    # 32 filters, 3x3 kernel, 'relu' activation
90    # 'input_shape' must match our image size (plus 3 for RGB color)
91    Conv2D(32, (3, 3), activation='relu', input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)),
92    MaxPooling2D(2, 2),
93
94    # Layer 2: Convolution + Pooling
95    Conv2D(64, (3, 3), activation='relu'),
96    MaxPooling2D(2, 2),
97
98    # Layer 3: Convolution + Pooling
99    Conv2D(128, (3, 3), activation='relu'),
100    MaxPooling2D(2, 2),
101
102    # Flatten the 3D feature maps into a 1D vector
103    Flatten(),
104
105    # A fully-connected 'Dense' layer with 512 neurons
106    Dense(512, activation='relu'),
107    
108    # 'Dropout' randomly turns off 50% of neurons during training
109    # This prevents the model from "memorizing" the training images (overfitting)
110    Dropout(0.5),
111
112    # The Output Layer: 1 neuron with 'sigmoid' activation
113    # 'sigmoid' outputs a single value between 0 and 1,
114    # perfect for binary (fresh/rotten) classification.
115    Dense(1, activation='sigmoid')
116])
117
118# --- 6. Compile the Model ---
119# Tell the model how to learn
120model.compile(
121    loss='binary_crossentropy', # The best loss function for binary (0 or 1) problems
122    optimizer='adam',           # A popular and effective optimizer
123    metrics=['accuracy']        # We want to see the 'accuracy' at each step
124)
125
126# Print a summary of our model architecture
127model.summary()
128
129# --- 7. Train the Model ---
130print("\nStarting model training...")
131history = model.fit(
132    train_generator,
133    steps_per_epoch=total_train // BATCH_SIZE, # Batches to run per epoch
134    epochs=EPOCHS,                             # Number of times to repeat
135    validation_data=validation_generator,
136    validation_steps=total_val // BATCH_SIZE,  # Batches to use for validation
137    verbose=1
138)
139print("Training complete!")
140
141# --- 8. Save the Trained Model ---
142# We save it in the modern '.keras' format
143model_filename = 'food_freshness_model.keras'
144model.save(model_filename)
145print(f"\nModel saved successfully as {model_filename}")
146
147# --- 9. Visualize Training Results ---
148# Plot accuracy and loss to see how the model performed
149acc = history.history['accuracy']
150val_acc = history.history['val_accuracy']
151loss = history.history['loss']
152val_loss = history.history['val_loss']
153
154epochs_range = range(EPOCHS)
155
156plt.figure(figsize=(12, 6))
157plt.subplot(1, 2, 1)
158plt.plot(epochs_range, acc, label='Training Accuracy')
159plt.plot(epochs_range, val_acc, label='Validation Accuracy')
160plt.legend(loc='lower right')
161plt.title('Training and Validation Accuracy')
162
163plt.subplot(1, 2, 2)
164plt.plot(epochs_range, loss, label='Training Loss')
165plt.plot(epochs_range, val_loss, label='Validation Loss')
166plt.legend(loc='upper right')
167plt.title('Training and Validation Loss')
168plt.show()