Enable training function with dataset/train-test and optional pseudo-labeling

README.md

@@ -37,10 +37,14 @@ Versi ini difokuskan untuk tahap awal:
   - hasil gambar deteksi
   - ringkasan klasifikasi
   - grafik jumlah person terdeteksi
+- **Menu Training** (accordion)
+  - membaca dataset dari `dataset/train` dan `dataset/test`
+  - jika label `.txt` YOLO belum ada, bisa aktifkan pseudo-label
+  - output model terbaik disimpan sebagai `best.pt`
 
 Catatan:
 - input evaluasi batch dihapus.
-- menu training dipisahkan dari halaman ini.
+- training tersedia sebagai menu terpisah (accordion) di halaman yang sama.
 - input jumlah presensi belum digunakan pada versi ini.
 
 ## File Penting

app.py

@@ -1,6 +1,9 @@
 import io
 import os
+import shutil
+from datetime import datetime
 from functools import lru_cache
+from pathlib import Path
 from typing import Any, Tuple
 
 import cv2
@@ -128,6 +131,161 @@ def run_check_and_classify(
         return image, summary, blank
 
 
+def _iter_images(folder: Path) -> list[Path]:
+    exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
+    if not folder.exists():
+        return []
+    return sorted([p for p in folder.rglob("*") if p.is_file() and p.suffix.lower() in exts])
+
+
+def _label_candidates(img_path: Path) -> list[Path]:
+    stem = img_path.stem
+    return [
+        img_path.with_suffix(".txt"),
+        img_path.parent.parent / "labels" / img_path.parent.name / f"{stem}.txt",
+        img_path.parent / "labels" / f"{stem}.txt",
+    ]
+
+
+def _write_pseudo_label(img_rgb: np.ndarray, label_path: Path, conf_threshold: float) -> int:
+    model = load_model()
+    h, w = img_rgb.shape[:2]
+    bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+    result = model.predict(source=bgr, conf=conf_threshold, verbose=False)[0]
+    person_id = _resolve_person_class_id(result)
+    boxes_xyxy = result.boxes.xyxy.cpu().numpy() if result.boxes is not None else np.empty((0, 4))
+    classes = result.boxes.cls.cpu().numpy() if result.boxes is not None else np.array([])
+
+    lines = []
+    for box, cls_id in zip(boxes_xyxy, classes):
+        if int(cls_id) != person_id:
+            continue
+        x1, y1, x2, y2 = box.tolist()
+        xc = ((x1 + x2) / 2) / w
+        yc = ((y1 + y2) / 2) / h
+        bw = (x2 - x1) / w
+        bh = (y2 - y1) / h
+        lines.append(f"0 {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}")
+
+    label_path.parent.mkdir(parents=True, exist_ok=True)
+    label_path.write_text("\n".join(lines), encoding="utf-8")
+    return len(lines)
+
+
+def _prepare_training_dataset(dataset_root: Path, use_pseudo_if_missing: bool, pseudo_conf: float) -> Tuple[Path, str]:
+    train_src = dataset_root / "train"
+    test_src = dataset_root / "test"
+    if not train_src.exists() or not test_src.exists():
+        raise FileNotFoundError("Folder dataset wajib berisi `dataset/train` dan `dataset/test`.")
+
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    prepared_root = dataset_root.parent / "prepared_dataset" / ts
+    img_train = prepared_root / "images" / "train"
+    img_val = prepared_root / "images" / "val"
+    lbl_train = prepared_root / "labels" / "train"
+    lbl_val = prepared_root / "labels" / "val"
+    for p in (img_train, img_val, lbl_train, lbl_val):
+        p.mkdir(parents=True, exist_ok=True)
+
+    def process_split(src: Path, out_img: Path, out_lbl: Path) -> Tuple[int, int]:
+        images = _iter_images(src)
+        if not images:
+            raise ValueError(f"Tidak ada gambar di {src}")
+        copied = 0
+        labeled = 0
+        for img in images:
+            target_img = out_img / img.name
+            if target_img.exists():
+                target_img = out_img / f"{img.stem}_{copied}{img.suffix}"
+            shutil.copy2(img, target_img)
+
+            label_target = out_lbl / f"{target_img.stem}.txt"
+            found_label = None
+            for cand in _label_candidates(img):
+                if cand.exists():
+                    found_label = cand
+                    break
+
+            if found_label is not None:
+                shutil.copy2(found_label, label_target)
+                labeled += 1
+            elif use_pseudo_if_missing:
+                rgb = cv2.cvtColor(cv2.imread(str(target_img)), cv2.COLOR_BGR2RGB)
+                labeled += 1 if _write_pseudo_label(rgb, label_target, pseudo_conf) >= 0 else 0
+            copied += 1
+        return copied, labeled
+
+    c_train, l_train = process_split(train_src, img_train, lbl_train)
+    c_val, l_val = process_split(test_src, img_val, lbl_val)
+    if l_train == 0 or l_val == 0:
+        raise ValueError(
+            "Label tidak ditemukan. Aktifkan pseudo-label atau sediakan file .txt YOLO untuk train/test."
+        )
+
+    yaml_path = prepared_root / "data.yaml"
+    yaml_path.write_text(
+        f"path: {prepared_root.as_posix()}\n"
+        "train: images/train\n"
+        "val: images/val\n"
+        "names:\n"
+        "  0: person\n",
+        encoding="utf-8",
+    )
+    info = (
+        f"Dataset siap: train={c_train} (labeled={l_train}), "
+        f"val={c_val} (labeled={l_val}), yaml={yaml_path}"
+    )
+    return yaml_path, info
+
+
+def run_training(
+    epochs: float,
+    imgsz: float,
+    batch: float,
+    use_pseudo_if_missing: bool,
+    pseudo_conf: float,
+) -> Tuple[str, str]:
+    try:
+        base_dir = Path(__file__).resolve().parent
+        dataset_root = base_dir / "dataset"
+        data_yaml, prep_info = _prepare_training_dataset(
+            dataset_root=dataset_root,
+            use_pseudo_if_missing=use_pseudo_if_missing,
+            pseudo_conf=float(pseudo_conf),
+        )
+
+        model_path = str(base_dir / "yolov8n.pt") if (base_dir / "yolov8n.pt").exists() else resolve_model_path()
+        model = YOLO(model_path)
+        run_name = f"classroom_train_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        results = model.train(
+            data=str(data_yaml),
+            epochs=int(epochs),
+            imgsz=int(imgsz),
+            batch=int(batch),
+            device="cpu",
+            project=str(base_dir / "runs_train"),
+            name=run_name,
+            exist_ok=False,
+            verbose=True,
+        )
+
+        save_dir = Path(str(results.save_dir))
+        best_pt = save_dir / "weights" / "best.pt"
+        best_out = base_dir / "best.pt"
+        if best_pt.exists():
+            shutil.copy2(best_pt, best_out)
+
+        msg = (
+            "Training selesai.\n"
+            f"- {prep_info}\n"
+            f"- Run dir: {save_dir}\n"
+            f"- best.pt: {best_out if best_out.exists() else best_pt}"
+        )
+        return msg, str(best_out) if best_out.exists() else ""
+    except Exception as err:
+        return f"Training gagal: {err}", ""
+
+
 CUSTOM_CSS = """
 .container {max-width: 1120px !important; margin: auto;}
 .section-box {padding: 14px; border: 1px solid #d7dee8; border-radius: 10px; background: #ffffff;}
@@ -161,12 +319,43 @@ with gr.Blocks(css=CUSTOM_CSS, title="Pengecekan dan Klasifikasi Mahasiswa - YOL
         class_result_md = gr.Markdown("Hasil klasifikasi akan tampil di sini.")
         comparison_chart = gr.Image(label="Grafik Klasifikasi", type="numpy")
 
+    with gr.Column(elem_classes=["section-box"]):
+        with gr.Accordion("Menu Training", open=False):
+            gr.Markdown(
+                "Training akan membaca dataset dari `dataset/train` dan `dataset/test`.\n"
+                "Jika label `.txt` belum ada, aktifkan pseudo-label agar proses tetap jalan."
+            )
+            with gr.Row():
+                train_epochs = gr.Number(label="Epochs", value=10, precision=0, minimum=1)
+                train_imgsz = gr.Number(label="Image Size", value=640, precision=0, minimum=64)
+                train_batch = gr.Number(label="Batch", value=8, precision=0, minimum=1)
+            use_pseudo = gr.Checkbox(
+                label="Gunakan pseudo-label jika label YOLO tidak tersedia",
+                value=True,
+            )
+            pseudo_conf = gr.Slider(
+                0.1,
+                0.9,
+                value=0.30,
+                step=0.05,
+                label="Confidence pseudo-label",
+            )
+            train_btn = gr.Button("Mulai Training", variant="primary")
+            train_log = gr.Textbox(label="Log Training", lines=8, interactive=False)
+            best_model_file = gr.File(label="best.pt hasil training")
+
     run_btn.click(
         fn=run_check_and_classify,
         inputs=[image_input, conf_input],
         outputs=[output_image, class_result_md, comparison_chart],
     )
 
+    train_btn.click(
+        fn=run_training,
+        inputs=[train_epochs, train_imgsz, train_batch, use_pseudo, pseudo_conf],
+        outputs=[train_log, best_model_file],
+    )
+
 
 if __name__ == "__main__":
     demo.launch()