Spaces:
Build error
Build error
File size: 4,389 Bytes
d869bdf fd11807 5bc80c7 fd11807 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Suppress TensorFlow logging
import cv2 as cv
import numpy as np
from tensorflow import keras
from keras import datasets, models, layers
from keras.callbacks import Callback
from keras.layers import Input
from ultralytics import YOLO # Assuming you have a YOLO.py file that contains a YOLO class
from tqdm import tqdm
# Suppress TensorFlow warnings and info messages
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
class TqdmProgressCallback(Callback):
def on_train_begin(self, logs=None):
self.total_batches = self.params['epochs'] * self.params['steps']
self.tqdm_bar = tqdm(total=self.total_batches, desc="Training Progress")
def on_batch_end(self, batch, logs=None):
self.tqdm_bar.update(1)
def on_train_end(self, logs=None):
self.tqdm_bar.close()
def resize_for_cnn(frame, target_size=(32, 32)):
return cv.resize(frame, target_size)
def preprocess_for_yolo(frame):
return cv.resize(frame, (620, 620)) # Assuming YOLOv8 input size is 640x640
def start():
try:
training_count = int(input("\nEnter the number of times you want to train the model: "))
except ValueError:
print("Please enter a valid integer for the training count.")
return
(training_images, training_labels), (testing_images, testing_labels) = datasets.cifar10.load_data()
training_images, testing_images = training_images / 255.0, testing_images / 255.0
input_shape = (32, 32, 3)
input_layer = Input(shape=input_shape)
conv1 = layers.Conv2D(32, (3, 3), activation='relu')(input_layer)
pool1 = layers.MaxPooling2D((2, 2))(conv1)
conv2 = layers.Conv2D(64, (3, 3), activation='relu')(pool1)
pool2 = layers.MaxPooling2D((2, 2))(conv2)
conv3 = layers.Conv2D(64, (3, 3), activation='relu')(pool2)
flatten = layers.Flatten()(conv3)
dense1 = layers.Dense(64, activation='relu')(flatten)
dense2 = layers.Dense(128, activation='relu', name='dense_layer_2')(dense1)
output_layer = layers.Dense(10, activation='softmax')(dense2)
model = models.Model(inputs=input_layer, outputs=output_layer)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
print("______________________________________________________________________________")
progress_callback = TqdmProgressCallback()
model.fit(training_images, training_labels, epochs=training_count, validation_data=(testing_images, testing_labels), verbose=0, callbacks=[progress_callback])
print("______________________________________________________________________________")
model.save('image_classifier.keras')
yolov8_model = YOLO('yolov8n-seg.pt')
cap = cv.VideoCapture(0)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
print("Error: Failed to capture frame.")
break
preprocessed_frame_yolo = preprocess_for_yolo(frame)
preprocessed_frame_cnn = resize_for_cnn(frame)
prediction = model.predict(np.expand_dims(preprocessed_frame_cnn, axis=0))
predicted_class = np.argmax(prediction)
detected_objects = yolov8_model(preprocessed_frame_yolo)
print("______________________________________________________________________________")
cv.imshow('Frame', frame)
if cv.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv.destroyAllWindows()
def end():
print("\nThe program has ended.")
def main():
while True:
command = input("Enter 'start' to start the program or 'end' to end the program: ").lower()
if command == 'start':
start()
break # Exit the loop after starting the program
elif command == 'end':
end()
break # Exit the loop after ending the program
else:
print("\nInvalid command. Please enter either 'start' or 'end'.")
print("______________________________________________________________________________")
if __name__ == "__main__":
print("______________________________________________________________________________")
main()
# Contributors:
# 1) Harith Kavish S
# 2) Sharwan Krishnan P
#
|