michaelflppv commited on
Commit
bfcc507
·
verified ·
1 Parent(s): 394d302

Upload 3 files

Browse files
Files changed (3) hide show
  1. LICENSE +21 -0
  2. main.py +163 -0
  3. requirements.txt +10 -0
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Mikhail Filippov
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
main.py ADDED
@@ -0,0 +1,163 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import tensorflow as tf
3
+ from tensorflow import keras
4
+ from keras import layers
5
+ from tensorflow.keras.preprocessing.image import ImageDataGenerator
6
+ import numpy as np
7
+ from PIL import Image
8
+
9
+ def create_model(input_shape=(32, 32, 3)):
10
+ """Create and return a CNN model for binary image classification."""
11
+ model = keras.Sequential([
12
+ layers.Input(shape=input_shape), # Proper input layer specification
13
+ layers.Conv2D(32, (3, 3), activation='relu'),
14
+ layers.MaxPooling2D((2, 2)),
15
+ layers.Conv2D(64, (3, 3), activation='relu'),
16
+ layers.MaxPooling2D((2, 2)),
17
+ layers.Conv2D(128, (3, 3), activation='relu'),
18
+ layers.MaxPooling2D((2, 2)),
19
+ layers.Flatten(),
20
+ layers.Dense(128, activation='relu'),
21
+ layers.Dense(1, activation='sigmoid')
22
+ ])
23
+
24
+ # Compile the model
25
+ model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
26
+ return model
27
+
28
+ def train_model(batch_size=32, epochs=8):
29
+ """Train the model and save it."""
30
+ # Generate data for training and validation
31
+ datagen = ImageDataGenerator(
32
+ rescale=1.0 / 255,
33
+ validation_split=0.2,
34
+ rotation_range=20, # Add data augmentation
35
+ width_shift_range=0.2,
36
+ height_shift_range=0.2,
37
+ shear_range=0.2,
38
+ zoom_range=0.2,
39
+ horizontal_flip=True
40
+ )
41
+
42
+ train_generator = datagen.flow_from_directory(
43
+ directory='archive/train',
44
+ target_size=(32, 32),
45
+ batch_size=batch_size,
46
+ class_mode='binary',
47
+ subset='training'
48
+ )
49
+
50
+ validation_generator = datagen.flow_from_directory(
51
+ directory='archive/train',
52
+ target_size=(32, 32),
53
+ batch_size=batch_size,
54
+ class_mode='binary',
55
+ subset='validation'
56
+ )
57
+
58
+ # Create model
59
+ model = create_model()
60
+
61
+ # Add early stopping to prevent overfitting
62
+ early_stopping = keras.callbacks.EarlyStopping(
63
+ monitor='val_loss',
64
+ patience=3,
65
+ restore_best_weights=True
66
+ )
67
+
68
+ # Train the model
69
+ history = model.fit(
70
+ train_generator,
71
+ validation_data=validation_generator,
72
+ epochs=epochs,
73
+ callbacks=[early_stopping]
74
+ )
75
+
76
+ # Evaluate the model
77
+ test_loss, test_acc = model.evaluate(validation_generator)
78
+ print(f'Test accuracy: {test_acc:.4f}')
79
+
80
+ # Save the model
81
+ model.save('trained_model.keras')
82
+ print("Model saved as 'trained_model.keras'")
83
+
84
+ return model, history
85
+
86
+ def load_and_preprocess_image(image_path, target_size=(32, 32)):
87
+ """Load and preprocess an image for prediction."""
88
+ try:
89
+ img = Image.open(image_path)
90
+ img = img.resize(target_size)
91
+ img = img.convert('RGB')
92
+ img_array = np.array(img) / 255.0
93
+ return np.expand_dims(img_array, axis=0)
94
+ except Exception as e:
95
+ print(f"Error processing image: {e}")
96
+ return None
97
+
98
+ def test_model(model_path='trained_model.keras'):
99
+ """Load a trained model and use it to classify an image."""
100
+ try:
101
+ # Load the trained model
102
+ model = tf.keras.models.load_model(model_path)
103
+ except Exception as e:
104
+ print(f"Error loading model: {e}")
105
+ return
106
+
107
+ # Path to the image to test
108
+ image_path = input('Enter the path to the image you want to test: ')
109
+
110
+ if not os.path.isfile(image_path):
111
+ print("Invalid path, please enter a valid path to an image.")
112
+ return
113
+
114
+ # Load and preprocess the image
115
+ input_image = load_and_preprocess_image(image_path)
116
+ if input_image is None:
117
+ return
118
+
119
+ # Predict the class of the image
120
+ prediction = model.predict(input_image, verbose=0)
121
+
122
+ # Define the threshold for classification
123
+ threshold = 0.5
124
+
125
+ # Classify the image
126
+ classification = "REAL" if prediction[0][0] > threshold else "FAKE"
127
+ confidence = prediction[0][0] if prediction[0][0] > threshold else 1 - prediction[0][0]
128
+
129
+ # Print the result
130
+ print(f"Classification: {classification}")
131
+ print(f"Confidence: {confidence * 100:.2f}%")
132
+ print(f"Raw prediction value: {prediction[0][0]:.4f}")
133
+
134
+ def main():
135
+ """Main function to run the program."""
136
+ # Set memory growth to avoid memory allocation errors
137
+ gpus = tf.config.experimental.list_physical_devices('GPU')
138
+ if gpus:
139
+ try:
140
+ for gpu in gpus:
141
+ tf.config.experimental.set_memory_growth(gpu, True)
142
+ except RuntimeError as e:
143
+ print(f"Error setting memory growth: {e}")
144
+
145
+ # Define hyperparameters
146
+ batch_size = 32
147
+ epochs = 10
148
+
149
+ while True:
150
+ activation_mode = input('Select mode (train/test/exit): ').lower()
151
+
152
+ if activation_mode == 'train':
153
+ train_model(batch_size, epochs)
154
+ elif activation_mode == 'test':
155
+ test_model()
156
+ elif activation_mode == 'exit':
157
+ print("Exiting program.")
158
+ break
159
+ else:
160
+ print('Invalid mode, please select "train", "test", or "exit"')
161
+
162
+ if __name__ == "__main__":
163
+ main()
requirements.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ # TensorFlow for deep learning
2
+ tensorflow==2.19.0
3
+ # NumPy for numerical operations
4
+ numpy==2.1.3
5
+ # Keras for building neural network models
6
+ keras==3.10.0
7
+ # Pillow for image processing
8
+ pillow==11.2.1
9
+ # SciPy for scientific computing
10
+ scipy==1.11.4