Create utils/split_data.py

utils/split_data.py

@@ -0,0 +1,76 @@
+import os
+import pandas as pd 
+import shutil 
+from sklearn.model_selection import train_test_split
+from tqdm.auto import tqdm
+
+df_2014 = pd.read_csv('data/CROHME/2014/caption.txt', sep='\t', header=None, names=['filenames', 'captions'])
+df_2016 = pd.read_csv('data/CROHME/2016/caption.txt', sep='\t', header=None, names=['filenames', 'captions'])
+df_2019 = pd.read_csv('data/CROHME/2019/caption.txt', sep='\t', header=None, names=['filenames', 'captions'])
+df_train = pd.read_csv('data/CROHME/train/caption.txt', sep='\t', header=None, names=['filenames', 'captions'])
+
+data = pd.concat(
+    [
+        df_2014,
+        df_2016,
+        df_2019,
+        df_train
+    ]
+)
+
+# First, split off 10% of the data to train and test sets
+train, test = train_test_split(data, test_size=0.1, random_state=42)
+
+# Second, split off 10% of the training data to train and validation sets 
+train, val = train_test_split(train, test_size=0.1, random_state=42)
+
+print("Train shape:", train.shape)
+print("Test shape:", test.shape)
+print("Validation shape:", val.shape)
+
+train_filenames = train['filenames'].tolist()
+train_captions = train['captions'].tolist()
+
+test_filenames = test['filenames'].tolist()
+test_captions = test['captions'].tolist()
+
+val_filenames = val['filenames'].tolist()
+val_captions = val['captions'].tolist()
+
+# Extract captions.txt for each split  
+with open('data/CROHME_splitted/train/caption.txt', 'w', encoding='utf-8') as f:
+    for filename, caption in zip(train_filenames, train_captions):
+        f.write(f"{filename}\t{caption}\n")
+
+with open('data/CROHME_splitted/test/caption.txt', 'w', encoding='utf-8') as f:
+    for filename, caption in zip(test_filenames, test_captions):
+        f.write(f"{filename}\t{caption}\n")
+
+with open('data/CROHME_splitted/val/caption.txt', 'w', encoding='utf-8') as f:
+    for filename, caption in zip(val_filenames, val_captions):
+        f.write(f"{filename}\t{caption}\n")
+
+
+IMAGES_DIR = 'data/images'
+TRAIN_DIR = 'data/CROHME_splitted/train/img'
+TEST_DIR = 'data/CROHME_splitted/test/img'
+VAL_DIR = 'data/CROHME_splitted/val/img'
+
+os.makedirs(TRAIN_DIR, exist_ok=True)
+os.makedirs(TEST_DIR, exist_ok=True)        
+os.makedirs(VAL_DIR, exist_ok=True)
+
+for train_filename in tqdm(train_filenames, desc="Copying train images"):
+    src = os.path.join(IMAGES_DIR, train_filename) + '.bmp'  # Ensure the file extension is correct
+    dst = os.path.join(TRAIN_DIR, train_filename) + '.bmp'
+    shutil.copy(src, dst)
+    
+for test_filename in tqdm(test_filenames, desc="Copying test images"):
+    src = os.path.join(IMAGES_DIR, test_filename) + '.bmp'
+    dst = os.path.join(TEST_DIR, test_filename) + '.bmp'
+    shutil.copy(src, dst)
+
+for val_filename in tqdm(val_filenames, desc="Copying validation images"):
+    src = os.path.join(IMAGES_DIR, val_filename) + '.bmp'
+    dst = os.path.join(VAL_DIR, val_filename) + '.bmp'
+    shutil.copy(src, dst)