Update burrito.py

burrito.py

@@ -0,0 +1,84 @@
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
+import numpy as np
+from tensorflow.keras.preprocessing import image
+
+# Define image size and batch size
+IMG_SIZE = 224
+BATCH_SIZE = 32
+
+# Define train and validation directories
+train_dir = 'path/to/train/folder'
+val_dir = 'path/to/validation/folder'
+
+# Use ImageDataGenerator for data augmentation
+train_datagen = ImageDataGenerator(
+    rescale=1./255,
+    rotation_range=20,
+    width_shift_range=0.1,
+    height_shift_range=0.1,
+    shear_range=0.2,
+    zoom_range=0.2,
+    horizontal_flip=True,
+    fill_mode='nearest')
+
+val_datagen = ImageDataGenerator(rescale=1./255)
+
+# Generate batches of augmented data from directories
+train_generator = train_datagen.flow_from_directory(
+    train_dir,
+    target_size=(IMG_SIZE, IMG_SIZE),
+    batch_size=BATCH_SIZE,
+    class_mode='categorical')
+
+val_generator = val_datagen.flow_from_directory(
+    val_dir,
+    target_size=(IMG_SIZE, IMG_SIZE),
+    batch_size=BATCH_SIZE,
+    class_mode='categorical')
+
+# Define the model
+model = Sequential()
+
+# Add convolutional layers
+model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(IMG_SIZE, IMG_SIZE, 3)))
+model.add(MaxPooling2D((2, 2)))
+model.add(Conv2D(64, (3, 3), activation='relu'))
+model.add(MaxPooling2D((2, 2)))
+model.add(Conv2D(128, (3, 3), activation='relu'))
+model.add(MaxPooling2D((2, 2)))
+
+# Flatten and add dense layers
+model.add(Flatten())
+model.add(Dense(512, activation='relu'))
+model.add(Dropout(0.5))
+model.add(Dense(3, activation='softmax'))
+
+# Compile the model
+model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
+
+# Train the model
+history = model.fit(
+    train_generator,
+    steps_per_epoch=train_generator.samples // BATCH_SIZE,
+    epochs=50,
+    validation_data=val_generator,
+    validation_steps=val_generator.samples // BATCH_SIZE)
+
+# Load an image
+img_path = 'path/to/image.jpg'
+img = image.load_img(img_path, target_size=(IMG_SIZE, IMG_SIZE))
+
+# Convert the image to a numpy array and normalize it
+img_array = image.img_to_array(img) / 255.
+
+# Expand the dimensions of the array to match the model input shape
+img_array = np.expand_dims(img_array, axis=0)
+
+# Use the model to predict the class probabilities
+probs = model.predict(img_array)[0]
+
+# Print the predicted class probabilities
+print(probs)