Initial commit

training/generate_shapes.py

@@ -0,0 +1,65 @@
+import os
+import cv2
+import numpy as np
+import random
+
+IMG_SIZE = 64
+NUM_IMAGES = 250  # Number of images per class
+OUTPUT_DIR = 'data'
+
+# Ensure folders exist
+shapes = ['circle', 'square', 'rectangle', 'triangle']
+for shape in shapes:
+    os.makedirs(os.path.join(OUTPUT_DIR, shape), exist_ok=True)
+
+def draw_circle():
+    img = np.ones((IMG_SIZE, IMG_SIZE), dtype=np.uint8) * 255
+    center = (random.randint(8, 56), random.randint(8, 56))
+    radius = random.randint(4, min(60 - max(center), min(center)))
+    cv2.circle(img, center, radius, (0,), 2)
+    return img
+
+def draw_square():
+    img = np.ones((IMG_SIZE, IMG_SIZE), dtype=np.uint8) * 255
+    start = (random.randint(4, 56), random.randint(4, 56))
+    size = random.randint(4, 60 - max(start))
+    cv2.rectangle(img, start, (start[0]+size + random.randint(0, 3), start[1]+size + random.randint(0, 3)), (0,), 2)
+    return img
+
+def draw_rectangle():
+    img = np.ones((IMG_SIZE, IMG_SIZE), dtype=np.uint8) * 255
+    vertical = random.randint(0, 1)
+    long, short = (random.randint(4, 28), random.randint(33, 56))
+    if vertical:
+        start = (random.randint(4, 60 - short), random.randint(4, 60 - long))
+        width, height = short, long
+    else:
+        start = (random.randint(4, 60 - long), random.randint(4, 60 - short))
+        width, height = long, short
+    cv2.rectangle(img, start, (start[0]+width, start[1]+height), (0,), 2)
+    return img
+
+def draw_triangle():
+    img = np.ones((IMG_SIZE, IMG_SIZE), dtype=np.uint8) * 255
+    pt1 = (random.randint(4, 60), random.randint(4, 60))
+    pt2 = (random.randint(4, 60), random.randint(4, 60))
+    pt3 = (random.randint(4, 60), random.randint(4, 60))
+    points = np.array([pt1, pt2, pt3])
+    cv2.drawContours(img, [points], 0, (0,), 2)
+    return img
+
+draw_functions = {
+    'circle': draw_circle,
+    'square': draw_square,
+    'rectangle': draw_rectangle,
+    'triangle': draw_triangle
+}
+
+# ----- Generate images -----
+for shape in shapes:
+    for i in range(NUM_IMAGES):
+        img = draw_functions[shape]()
+        filename = os.path.join(OUTPUT_DIR, shape, f"{shape}_{i}.png")
+        cv2.imwrite(filename, img)
+
+print("Images generated successfully!")

training/train.py

@@ -0,0 +1,68 @@
+import tensorflow as tf
+from keras import layers, models
+from keras.preprocessing import image
+import pathlib
+import numpy as np
+
+IMG_SIZE = 64
+BATCH_SIZE = 8
+NUM_CLASSES = 4
+EPOCHS = 10
+
+data_dir = pathlib.Path("data")
+
+train_ds = tf.keras.utils.image_dataset_from_directory(
+    data_dir,
+    labels='inferred',
+    label_mode='categorical',
+    color_mode='grayscale',
+    batch_size=BATCH_SIZE,
+    image_size=(IMG_SIZE, IMG_SIZE),
+    validation_split=0.2,
+    subset="training",
+    seed=123
+)
+
+val_ds = tf.keras.utils.image_dataset_from_directory(
+    data_dir,
+    labels='inferred',
+    label_mode='categorical',
+    color_mode='grayscale',
+    batch_size=BATCH_SIZE,
+    image_size=(IMG_SIZE, IMG_SIZE),
+    validation_split=0.2,
+    subset="validation",
+    seed=123
+)
+
+AUTOTUNE = tf.data.AUTOTUNE
+train_ds = train_ds.cache().shuffle(100).prefetch(buffer_size=AUTOTUNE)
+val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
+
+model = models.Sequential([
+    layers.Rescaling(1/255, input_shape=(IMG_SIZE, IMG_SIZE, 1)),
+    layers.Conv2D(32, (3,3), activation='relu'),
+    layers.MaxPooling2D(2,2),
+
+    layers.Conv2D(64, (3,3), activation='relu'),
+    layers.MaxPooling2D(2,2),
+
+    layers.Flatten(),
+    layers.Dense(64, activation='relu'),
+    layers.Dense(NUM_CLASSES, activation='softmax')
+])
+
+model.compile(
+    optimizer='adam',
+    loss='categorical_crossentropy',
+    metrics=['accuracy']
+)
+
+model.fit(
+    train_ds,
+    validation_data=val_ds,
+    epochs=EPOCHS
+)
+
+model.save("shape_model.keras")
+print("Saved model")