Update augmentation script to scale Teacher dataset to 900 images

Scaled the base dataset target from 100 to 300 images per class (Blast, Brown Spot, Sheath Blight) to resolve underfitting and poor accuracy (35%) in the initial Teacher model.
Retained geometric transformations (horizontal/vertical flips) with precise YOLO polygon coordinate inversions.
Retained photometric transformations (Gaussian noise and brightness scaling) to improve model robustness against varied environmental drone lighting.

augment.py

@@ -6,143 +6,130 @@ from glob import glob
 from tqdm import tqdm
 
 # ================= CONFIGURATION =================
-# 1. Path to your Seed Dataset (The 36 images you just labelled)
-IMAGES_DIR = r"C:\Users\charu\Desktop\My_Project\seed_images"
-LABELS_DIR = r"C:\Users\charu\Desktop\My_Project\seed_labels"
+# 1. PATHS (Separated)
+IMAGES_DIR = r"C:\Users\charu\Desktop\04-02-2026\images"
+LABELS_DIR = r"C:\Users\charu\Desktop\04-02-2026\labels"
 
-# 2. Target count per class (You wanted 100)
-TARGET_COUNT = 100 
+# 2. Target Count per class
+TARGET_PER_CLASS = 300
 
-# 3. Your Classes
-# Note: "Healthy" is identified by empty text files
-CLASS_NAMES = {0: "Blast", 1: "Brown Spot", 2: "Sheath Blight", "Healthy": "Healthy"}
+# 3. Class Names
+CLASS_NAMES = {0: "Blast", 1: "Brown Spot", 2: "Sheath Blight"}
 # =================================================
 
-def load_data():
-    """Sorts images into lists based on what disease they contain."""
-    data_map = {0: [], 1: [], 2: [], "Healthy": []}
+def load_dataset():
+    dataset = {0: [], 1: [], 2: []}
     
-    # Get all text files
+    # scan labels folder
     txt_files = glob(os.path.join(LABELS_DIR, "*.txt"))
     
+    print(f"📂 Scanning Labels: {LABELS_DIR}")
+    print(f"   -> Found {len(txt_files)} text files.")
+    
+    if len(txt_files) == 0:
+        print(" Error: No text files found! Check the path.")
+        return dataset
+
     for txt_path in txt_files:
         filename = os.path.basename(txt_path).replace('.txt', '')
         
-        # Find matching image
+        # Look for matching image in IMAGES_DIR
         img_jpg = os.path.join(IMAGES_DIR, filename + ".jpg")
         img_png = os.path.join(IMAGES_DIR, filename + ".png")
+        img_jpeg = os.path.join(IMAGES_DIR, filename + ".jpeg")
         
         if os.path.exists(img_jpg): img_path = img_jpg
         elif os.path.exists(img_png): img_path = img_png
-        else: continue # Skip if image missing
+        elif os.path.exists(img_jpeg): img_path = img_jpeg
+        else: 
+            # If no image found for this label, skip it
+            continue 
 
-        # Read the label
         with open(txt_path, 'r') as f:
             lines = f.readlines()
-        
-        # CLASSIFY THE IMAGE
-        if not lines:
-            # Empty file = Healthy
-            data_map["Healthy"].append((img_path, lines))
-        else:
-            # Check the first class ID in the file
-            class_id = int(lines[0].split()[0])
-            if class_id in data_map:
-                data_map[class_id].append((img_path, lines))
-    
-    return data_map
-
-def augment_and_save(img_path, lines, new_name):
+            
+        if lines:
+            try:
+                # Read class ID
+                class_id = int(lines[0].split()[0])
+                if class_id in dataset:
+                    dataset[class_id].append((img_path, lines))
+            except:
+                pass 
+    return dataset
+
+def augment_polygon(img_path, lines, new_filename):
     img = cv2.imread(img_path)
     if img is None: return
-
-    # Random Augmentation Strategy
-    aug_type = random.choice(["h_flip", "v_flip", "noise", "bright", "dark"])
+    
+    action = random.choice(["h_flip", "v_flip", "bright", "noise"])
     new_lines = []
     
-    h, w, _ = img.shape
-
-    # 1. Horizontal Flip
-    if aug_type == "h_flip":
+    if action == "h_flip":
         new_img = cv2.flip(img, 1)
         for line in lines:
             parts = line.strip().split()
             cls = parts[0]
             coords = [float(x) for x in parts[1:]]
-            # Math: New X = 1.0 - Old X
-            new_coords = [1.0 - val if i % 2 == 0 else val for i, val in enumerate(coords)]
+            new_coords = []
+            for i, val in enumerate(coords):
+                if i % 2 == 0: new_coords.append(1.0 - val) # X
+                else:          new_coords.append(val)       # Y
             new_lines.append(f"{cls} " + " ".join([f"{c:.6f}" for c in new_coords]) + "\n")
 
-    # 2. Vertical Flip
-    elif aug_type == "v_flip":
+    elif action == "v_flip":
         new_img = cv2.flip(img, 0)
         for line in lines:
             parts = line.strip().split()
             cls = parts[0]
             coords = [float(x) for x in parts[1:]]
-            # Math: New Y = 1.0 - Old Y
-            new_coords = [1.0 - val if i % 2 != 0 else val for i, val in enumerate(coords)]
+            new_coords = []
+            for i, val in enumerate(coords):
+                if i % 2 == 0: new_coords.append(val)       # X
+                else:          new_coords.append(1.0 - val) # Y
             new_lines.append(f"{cls} " + " ".join([f"{c:.6f}" for c in new_coords]) + "\n")
 
-    # 3. Noise (Grainy)
-    elif aug_type == "noise":
-        noise = np.random.normal(0, 25, img.shape).astype(np.uint8)
-        new_img = cv2.add(img, noise)
-        new_lines = lines # Coordinates don't change
-
-    # 4. Brightness
-    elif aug_type == "bright":
-        new_img = cv2.convertScaleAbs(img, alpha=1.1, beta=20)
-        new_lines = lines
+    elif action == "bright":
+        beta = random.randint(-30, 30)
+        new_img = cv2.convertScaleAbs(img, alpha=1.0, beta=beta)
+        new_lines = lines 
 
-    # 5. Darkness
-    elif aug_type == "dark":
-        new_img = cv2.convertScaleAbs(img, alpha=0.9, beta=-20)
+    elif action == "noise":
+        noise = np.random.normal(0, 15, img.shape).astype(np.uint8)
+        new_img = cv2.add(img, noise)
+        new_lines = lines 
+        
+    else: 
+        new_img = img
         new_lines = lines
 
-    # Save Image
-    cv2.imwrite(os.path.join(IMAGES_DIR, new_name + ".jpg"), new_img)
-    
-    # Save Label
-    with open(os.path.join(LABELS_DIR, new_name + ".txt"), 'w') as f:
+    # SAVE TO SEPARATE FOLDERS
+    cv2.imwrite(os.path.join(IMAGES_DIR, new_filename + ".jpg"), new_img)
+    with open(os.path.join(LABELS_DIR, new_filename + ".txt"), 'w') as f:
         f.writelines(new_lines)
 
 def main():
-    print("🚀 scanning seed dataset...")
-    data_map = load_data()
+    print("🚀 Loading Dataset (Separated Folders)...")
+    data_map = load_dataset()
     
-    print("\n📊 Current counts:")
-    for key, items in data_map.items():
-        name = CLASS_NAMES[key]
-        print(f"  - {name}: {len(items)} images")
-
-    print(f"\n🛠️ Augmenting to reach {TARGET_COUNT} per class...")
-
-    for key, items in data_map.items():
-        name = CLASS_NAMES[key]
-        current_count = len(items)
-        needed = TARGET_COUNT - current_count
+    print("\n📊 Current Counts:")
+    for cid in [0, 1, 2]:
+        print(f"   - {CLASS_NAMES[cid]}: {len(data_map[cid])} images")
         
-        if needed <= 0:
-            print(f"✅ {name} is already full. Skipping.")
-            continue
-            
-        if current_count == 0:
-            print(f"⚠️ Warning: No images found for {name}. Cannot augment!")
-            continue
-
-        print(f"  -> Generating {needed} images for {name}...")
+    print("\n🛠️ augmenting...")
+    for cid in [0, 1, 2]:
+        items = data_map[cid]
+        needed = TARGET_PER_CLASS - len(items)
         
-        for i in tqdm(range(needed)):
-            # Pick a random source image to clone
-            source_img, source_lines = random.choice(items)
-            
-            # Create unique name
-            new_filename = f"aug_{key}_{i}"
-            
-            augment_and_save(source_img, source_lines, new_filename)
-            
-    print("\n✅ Augmentation Complete! You now have 400 images.")
+        if needed > 0 and items:
+            print(f"   -> Generating {needed} images for {CLASS_NAMES[cid]}...")
+            for i in tqdm(range(needed)):
+                src_img, src_lines = random.choice(items)
+                augment_polygon(src_img, src_lines, f"aug_{cid}_{i}")
+        elif not items:
+            print(f"⚠️ Warning: No images found for {CLASS_NAMES[cid]}!")
+
+    print("\n✅ Done!")
 
 if __name__ == "__main__":
     main()