Merge branch 'main' of https://github.com/ATLAS-PJK-GM007/ZeaVis-Edu into selly/frontend
This commit is contained in:
seriouselly
@@ -62,20 +62,44 @@ jobs:
|
||||
echo "HF_TOKEN is set (length: ${#HF_TOKEN})"
|
||||
echo "::endgroup::"
|
||||
|
||||
echo "::group::Download model.onnx"
|
||||
echo "::group::Download model files from Hugging Face"
|
||||
python -c "
|
||||
from huggingface_hub import hf_hub_download
|
||||
import os
|
||||
os.makedirs('model', exist_ok=True)
|
||||
path = hf_hub_download(
|
||||
repo_id='MythEclipse2737/corn-leaf-disease-classifier',
|
||||
filename='model/model.onnx',
|
||||
token=os.environ['HF_TOKEN'],
|
||||
local_dir='.',
|
||||
)
|
||||
print(f'Downloaded: {path}')
|
||||
import os, shutil
|
||||
repo = 'MythEclipse2737/zeavis-edu-corn-leaf-classifier'
|
||||
token = os.environ['HF_TOKEN']
|
||||
base = os.path.abspath('.')
|
||||
|
||||
# Files sit at root of HF repo → copy to correct subdirs
|
||||
# model.onnx goes to model/ for Docker COPY
|
||||
os.makedirs(os.path.join(base, 'model'), exist_ok=True)
|
||||
os.makedirs(os.path.join(base, 'best_model'), exist_ok=True)
|
||||
|
||||
# ONNX → model/model.onnx (Docker expects this path)
|
||||
p = hf_hub_download(repo_id=repo, filename='model.onnx', token=token)
|
||||
shutil.copy2(p, os.path.join(base, 'model', 'model.onnx'))
|
||||
print('model/model.onnx OK')
|
||||
|
||||
# TFLite (optional, for edge)
|
||||
p = hf_hub_download(repo_id=repo, filename='model.tflite', token=token)
|
||||
shutil.copy2(p, os.path.join(base, 'model', 'model.tflite'))
|
||||
print('model/model.tflite OK')
|
||||
|
||||
# Labels
|
||||
p = hf_hub_download(repo_id=repo, filename='labels.json', token=token)
|
||||
shutil.copy2(p, os.path.join(base, 'model', 'labels.json'))
|
||||
print('model/labels.json OK')
|
||||
|
||||
# Keras model + calibration for re-export
|
||||
p = hf_hub_download(repo_id=repo, filename='best_model.keras', token=token)
|
||||
shutil.copy2(p, os.path.join(base, 'best_model', 'best_model.keras'))
|
||||
print('best_model/best_model.keras OK')
|
||||
|
||||
p = hf_hub_download(repo_id=repo, filename='calibration.json', token=token)
|
||||
shutil.copy2(p, os.path.join(base, 'best_model', 'calibration.json'))
|
||||
print('best_model/calibration.json OK')
|
||||
"
|
||||
ls -lh model/model.onnx
|
||||
ls -lh model/model.onnx model/model.tflite model/labels.json best_model/best_model.keras 2>/dev/null
|
||||
echo "::endgroup::"
|
||||
|
||||
- name: Log in to GHCR
|
||||
|
||||
@@ -16,4 +16,7 @@ dataset/
|
||||
dataset_split/
|
||||
dataset_jagung.zip
|
||||
best_model/
|
||||
model/
|
||||
model/
|
||||
dataset*.zip
|
||||
dataset_*/
|
||||
corn-leaf-disease.zip
|
||||
Binary file not shown.
@@ -30,7 +30,7 @@ def main():
|
||||
parser.add_argument(
|
||||
"--opset",
|
||||
type=int,
|
||||
default=13,
|
||||
default=18,
|
||||
help="ONNX opset version to target",
|
||||
)
|
||||
|
||||
|
||||
+1835
-2947
File diff suppressed because one or more lines are too long
@@ -2,6 +2,13 @@ import os
|
||||
import shutil
|
||||
import zipfile
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# ==========================================
|
||||
# KONFIGURASI DAN MAPPING
|
||||
@@ -9,17 +16,25 @@ import json
|
||||
DAFTAR_ZIP = ['dataset_1.zip', 'dataset_2.zip', 'dataset_3.zip']
|
||||
TARGET_DIR = "dataset"
|
||||
|
||||
# Kaggle dataset paths (untuk auto-download)
|
||||
KAGGLE_DS1 = "ndisan/corn-leaf-disease"
|
||||
KAGGLE_DS2 = "smaranjitghose/corn-or-maize-leaf-disease-dataset"
|
||||
|
||||
# Google Drive file ID for dataset_3 fallback (SciDB via alternative host)
|
||||
DS3_GDRIVE_ID = None # Replace with known ID if available
|
||||
|
||||
PEMETAAN_KATEGORI = {
|
||||
"大斑病": "Hawar Daun",
|
||||
"小斑病": "Hawar Daun",
|
||||
"褐斑病": "Bercak Daun",
|
||||
"弯孢霉叶斑病": "Bercak Daun",
|
||||
"圆斑病": "Bercak Daun",
|
||||
"灰斑病": "Bercak Daun",
|
||||
"南方锈病": "Karat Daun",
|
||||
"普通锈病": "Karat Daun",
|
||||
"大斑病": "Hawar Daun",
|
||||
"小斑病": "Hawar Daun",
|
||||
"褐斑病": "Bercak Daun",
|
||||
"弯孢霉叶斑病": "Bercak Daun",
|
||||
"圆斑病": "Bercak Daun",
|
||||
"灰斑病": "Bercak Daun",
|
||||
"南方锈病": "Karat Daun",
|
||||
"普通锈病": "Karat Daun",
|
||||
}
|
||||
|
||||
# Known-corrupt images (empty/broken JPEG headers)
|
||||
DAFTAR_FILE_HAPUS = [
|
||||
"CBS28.jpg",
|
||||
"Corn_Common_Rust (1275).jpg",
|
||||
@@ -28,11 +43,113 @@ DAFTAR_FILE_HAPUS = [
|
||||
"Corn_Gray_Spot (1).jpg"
|
||||
]
|
||||
|
||||
MIN_WIDTH, MIN_HEIGHT = 32, 32
|
||||
MAX_ASPECT_RATIO = 5.0
|
||||
MIN_FILE_SIZE = 512
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 0: DOWNLOAD DATASET OTOMATIS
|
||||
# ==========================================
|
||||
def _download_kagglehub(ref, zip_name, timeout=120):
|
||||
"""Download dataset via kagglehub with a timeout to prevent hanging."""
|
||||
import multiprocessing
|
||||
log.info(f" Downloading {ref} via kagglehub (timeout={timeout}s)...")
|
||||
|
||||
def _do_download(q):
|
||||
try:
|
||||
import kagglehub
|
||||
path = kagglehub.dataset_download(ref)
|
||||
q.put(("ok", path))
|
||||
except Exception as e:
|
||||
q.put(("err", str(e)))
|
||||
|
||||
q = multiprocessing.Queue()
|
||||
p = multiprocessing.Process(target=_do_download, args=(q,))
|
||||
p.start()
|
||||
p.join(timeout)
|
||||
|
||||
if p.is_alive():
|
||||
p.terminate()
|
||||
p.join()
|
||||
log.warning(f" [FAIL] {ref} download timed out after {timeout}s.")
|
||||
return None
|
||||
|
||||
status, val = q.get()
|
||||
if status != "ok":
|
||||
log.warning(f" [FAIL] {ref} download failed: {val}")
|
||||
return None
|
||||
|
||||
path = val
|
||||
# Pack into ZIP for extraction step
|
||||
try:
|
||||
shutil.make_archive(zip_name.replace('.zip', ''), 'zip', path)
|
||||
log.info(f" [OK] {ref} -> {zip_name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
log.warning(f" [FAIL] Could not pack {path} into {zip_name}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def download_dataset_1():
|
||||
"""Download dataset_1 (ndisan/corn-leaf-disease) from Kaggle."""
|
||||
try:
|
||||
import kagglehub
|
||||
except ImportError:
|
||||
log.warning(" [SKIP] kagglehub not installed. Install: pip install kagglehub")
|
||||
return None
|
||||
return _download_kagglehub("ndisan/corn-leaf-disease", "dataset_1.zip")
|
||||
|
||||
|
||||
def download_dataset_2():
|
||||
"""Download dataset_2 (smaranjitghose/corn-or-maize-leaf-disease-dataset) from Kaggle."""
|
||||
try:
|
||||
import kagglehub
|
||||
except ImportError:
|
||||
log.warning(" [SKIP] kagglehub not installed. Install: pip install kagglehub")
|
||||
return None
|
||||
return _download_kagglehub("smaranjitghose/corn-or-maize-leaf-disease-dataset", "dataset_2.zip")
|
||||
|
||||
|
||||
def download_dataset_3():
|
||||
"""Download dataset_3 (SciDB China Agricultural Dataset).
|
||||
|
||||
Attempt: Kaggle alternative host, then Google Drive, then warn.
|
||||
If all fail, user must download manually from SciDB.
|
||||
"""
|
||||
# Try Kaggle alternative (if available)
|
||||
KAGGLE_DS3 = "disease-identification/corn-leaf-disease-chinese" # not guaranteed
|
||||
try:
|
||||
import kagglehub
|
||||
path = kagglehub.dataset_download(KAGGLE_DS3)
|
||||
log.info(f" [OK] dataset_3 downloaded from Kaggle mirror to {path}")
|
||||
return path
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
log.warning(" [SKIP] dataset_3 could not be downloaded automatically.")
|
||||
log.warning(" Please download manually from:")
|
||||
log.warning(" https://www.scidb.cn/en/detail?dataSetId=19536c73f6d74946a212719a94f53ab3")
|
||||
return None
|
||||
|
||||
|
||||
def download_dataset(zip_name, download_fn):
|
||||
"""Attempt auto-download if ZIP doesn't exist locally."""
|
||||
if os.path.exists(zip_name):
|
||||
log.info(f" [CACHE] {zip_name} already exists, skipping download.")
|
||||
return True
|
||||
log.info(f"--- Downloading {zip_name} ---")
|
||||
result = download_fn()
|
||||
if not result: # None or False
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 1: EKSTRAKSI DATASET
|
||||
# ==========================================
|
||||
def ekstrak_semua_zip():
|
||||
print("--- TAHAP 1: Mengekstrak File ZIP ---")
|
||||
log.info("--- TAHAP 1: Mengekstrak File ZIP ---")
|
||||
for zip_file in DAFTAR_ZIP:
|
||||
if os.path.exists(zip_file):
|
||||
folder_name = os.path.splitext(zip_file)[0]
|
||||
@@ -40,50 +157,58 @@ def ekstrak_semua_zip():
|
||||
try:
|
||||
with zipfile.ZipFile(zip_file, 'r') as zip_ref:
|
||||
zip_ref.extractall(folder_name)
|
||||
print(f" [OK] {zip_file} -> {folder_name}/")
|
||||
log.info(f" [OK] {zip_file} -> {folder_name}/")
|
||||
except zipfile.BadZipFile:
|
||||
print(f" [ERROR] {zip_file} rusak.")
|
||||
log.error(f" [ERROR] {zip_file} rusak.")
|
||||
else:
|
||||
print(f" [SKIP] File {zip_file} tidak ditemukan.")
|
||||
print("\n")
|
||||
log.warning(f" [SKIP] File {zip_file} tidak ditemukan. Beberapa kelas mungkin kosong.")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 2: GABUNGKAN DATASET 1 & 2
|
||||
# ==========================================
|
||||
def cari_folder_ds2(base_path, keywords):
|
||||
if not os.path.exists(base_path): return None
|
||||
if not os.path.exists(base_path):
|
||||
return None
|
||||
for f in os.listdir(base_path):
|
||||
f_lower = f.lower()
|
||||
if any(k in f_lower for k in keywords):
|
||||
return os.path.join(base_path, f)
|
||||
return None
|
||||
|
||||
|
||||
def gabungkan_dataset_1_dan_2():
|
||||
print("--- TAHAP 2: Menggabungkan Dataset 1 & 2 ---")
|
||||
log.info("--- TAHAP 2: Menggabungkan Dataset 1 & 2 ---")
|
||||
os.makedirs(TARGET_DIR, exist_ok=True)
|
||||
|
||||
# 1. Salin dari dataset_1
|
||||
# 1. Salin dari dataset_1
|
||||
folder_dari_ds1 = ["Bercak Daun", "Daun Sehat", "Hawar Daun"]
|
||||
for folder in folder_dari_ds1:
|
||||
src = os.path.join("dataset_1", folder)
|
||||
dst = os.path.join(TARGET_DIR, folder)
|
||||
if os.path.exists(src):
|
||||
shutil.copytree(src, dst, dirs_exist_ok=True)
|
||||
print(f" [OK] Menyalin folder {src} ke {dst}")
|
||||
log.info(f" [OK] Menyalin folder {src} ke {dst}")
|
||||
|
||||
# 2. Salin gambar dari dataset_2
|
||||
base_ds2 = os.path.join("dataset_2", "data")
|
||||
mapping_ds2 = {
|
||||
("common_rust", "commont_rust"): "Karat Daun",
|
||||
("healthy",): "Daun Sehat"
|
||||
}
|
||||
# 2. Dataset 2 — cari subfolder yang cocok
|
||||
if os.path.exists("dataset_2"):
|
||||
# Cari folder data/ atau folder langsung
|
||||
base_ds2 = "dataset_2"
|
||||
# Check for nested structure (Kaggle download structure varies)
|
||||
data_sub = os.path.join(base_ds2, "data")
|
||||
if os.path.exists(data_sub):
|
||||
base_ds2 = data_sub
|
||||
|
||||
mapping_ds2 = {
|
||||
("common_rust", "commont_rust"): "Karat Daun",
|
||||
("healthy",): "Daun Sehat"
|
||||
}
|
||||
|
||||
if os.path.exists(base_ds2):
|
||||
for keywords, target_subfolder in mapping_ds2.items():
|
||||
src_folder = cari_folder_ds2(base_ds2, keywords)
|
||||
dst_folder = os.path.join(TARGET_DIR, target_subfolder)
|
||||
os.makedirs(dst_folder, exist_ok=True)
|
||||
|
||||
|
||||
if src_folder and os.path.exists(src_folder):
|
||||
file_count = 0
|
||||
for file_name in os.listdir(src_folder):
|
||||
@@ -91,10 +216,12 @@ def gabungkan_dataset_1_dan_2():
|
||||
if os.path.isfile(full_file_name):
|
||||
shutil.copy(full_file_name, dst_folder)
|
||||
file_count += 1
|
||||
print(f" [OK] Menyalin {file_count} gambar dari {src_folder} ke {dst_folder}")
|
||||
log.info(f" [OK] Menyalin {file_count} gambar dari {src_folder} ke {dst_folder}")
|
||||
else:
|
||||
print(f" [SKIP] Folder untuk '{target_subfolder}' tidak ditemukan di {base_ds2}")
|
||||
print("\n")
|
||||
log.warning(f" [SKIP] Folder untuk '{target_subfolder}' tidak ditemukan di dataset_2")
|
||||
else:
|
||||
log.warning(" [SKIP] Folder dataset_2/ tidak ada, dataset_2 tidak diproses.")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 3: GABUNGKAN DATASET 3 (JSON MAPPING)
|
||||
@@ -104,20 +231,22 @@ def cari_gambar_fleksibel(folder_sumber, nama_file_target):
|
||||
path_langsung = os.path.join(folder_sumber, nama_file_target)
|
||||
if os.path.exists(path_langsung):
|
||||
return path_langsung
|
||||
|
||||
|
||||
target_lower = nama_file_target.lower()
|
||||
for f in os.listdir(folder_sumber):
|
||||
if f.lower() == target_lower or os.path.splitext(f)[0].lower() == os.path.splitext(target_lower)[0]:
|
||||
if (f.lower() == target_lower
|
||||
or os.path.splitext(f)[0].lower() == os.path.splitext(target_lower)[0]):
|
||||
return os.path.join(folder_sumber, f)
|
||||
return None
|
||||
|
||||
|
||||
def gabungkan_dataset_3():
|
||||
print("--- TAHAP 3: Menggabungkan Dataset 3 berdasarkan JSON ---")
|
||||
log.info("--- TAHAP 3: Menggabungkan Dataset 3 berdasarkan JSON ---")
|
||||
folder_data = os.path.join("dataset_3", "data")
|
||||
file_json = os.path.join("dataset_3", "desc.json")
|
||||
|
||||
if not os.path.exists(file_json):
|
||||
print(f" [SKIP] {file_json} tidak ditemukan.\n")
|
||||
log.warning(f" [SKIP] {file_json} tidak ditemukan. Dataset 3 dilewati.")
|
||||
return
|
||||
|
||||
with open(file_json, 'r', encoding='utf-8') as f:
|
||||
@@ -139,13 +268,14 @@ def gabungkan_dataset_3():
|
||||
shutil.copy(path_sumber, os.path.join(folder_tujuan, nama_asli))
|
||||
berhasil += 1
|
||||
|
||||
print(f" [OK] Berhasil merutekan {berhasil} gambar dari dataset_3 ke '{TARGET_DIR}'\n")
|
||||
log.info(f" [OK] Berhasil merutekan {berhasil} gambar dari dataset_3 ke '{TARGET_DIR}'")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 4: PEMBERSIHAN DATA (CLEANING)
|
||||
# ==========================================
|
||||
def bersihkan_dataset():
|
||||
print("--- TAHAP 4: Menghapus File Spesifik ---")
|
||||
log.info("--- TAHAP 4: Menghapus File Bermasalah ---")
|
||||
set_hapus = set(DAFTAR_FILE_HAPUS)
|
||||
terhapus = 0
|
||||
|
||||
@@ -156,42 +286,219 @@ def bersihkan_dataset():
|
||||
path_lengkap = os.path.join(root, nama_file)
|
||||
try:
|
||||
os.remove(path_lengkap)
|
||||
print(f" [TERHAPUS] {path_lengkap}")
|
||||
set_hapus.remove(nama_file)
|
||||
log.info(f" [TERHAPUS] {path_lengkap}")
|
||||
set_hapus.discard(nama_file)
|
||||
terhapus += 1
|
||||
except Exception as e:
|
||||
print(f" [GAGAL] {path_lengkap} ({e})")
|
||||
|
||||
print(f" [OK] Total file dihapus: {terhapus}")
|
||||
log.warning(f" [GAGAL] {path_lengkap} ({e})")
|
||||
|
||||
log.info(f" [OK] Total file spesifik dihapus: {terhapus}")
|
||||
if set_hapus:
|
||||
print(f" [INFO] {len(set_hapus)} file tidak ditemukan (mungkin sudah terhapus sebelumnya):")
|
||||
log.info(f" [INFO] {len(set_hapus)} file tidak ditemukan (mungkin sudah terhapus sebelumnya)")
|
||||
for sisa in set_hapus:
|
||||
print(f" - {sisa}")
|
||||
print("\n")
|
||||
log.info(f" - {sisa}")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 5: BUNGKUS KE ZIP
|
||||
# TAHAP 5: HAPUS AUGMENTED DUPLICATES
|
||||
# ==========================================
|
||||
def hapus_augmented_duplicates():
|
||||
"""Remove pre-augmented duplicates (augmented_* files)."""
|
||||
log.info("--- TAHAP 5: Menghapus Augmented Duplicates ---")
|
||||
removed = 0
|
||||
if not os.path.exists(TARGET_DIR):
|
||||
log.warning(" [SKIP] Dataset folder tidak ditemukan.")
|
||||
return
|
||||
|
||||
for root, _, files in os.walk(TARGET_DIR):
|
||||
for nama_file in files:
|
||||
if nama_file.startswith("augmented_"):
|
||||
path_file = os.path.join(root, nama_file)
|
||||
try:
|
||||
os.remove(path_file)
|
||||
removed += 1
|
||||
except Exception as e:
|
||||
log.warning(f" [GAGAL] {path_file} ({e})")
|
||||
|
||||
log.info(f" [OK] {removed} augmented_* files dihapus.")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 6: CORRUPT & INVALID IMAGE DETECTION
|
||||
# ==========================================
|
||||
def validasi_gambar():
|
||||
log.info("--- TAHAP 6: Validasi & Deteksi Gambar Rusak ---")
|
||||
dihapus_total = 0
|
||||
stat_count = {
|
||||
"too_small_file": 0,
|
||||
"cannot_open": 0,
|
||||
"too_small_dims": 0,
|
||||
"extreme_aspect": 0,
|
||||
}
|
||||
|
||||
if not os.path.exists(TARGET_DIR):
|
||||
log.warning(" [SKIP] Dataset folder tidak ditemukan.")
|
||||
return
|
||||
|
||||
for root, _, files in os.walk(TARGET_DIR):
|
||||
for nama_file in files:
|
||||
path_file = os.path.join(root, nama_file)
|
||||
|
||||
try:
|
||||
file_size = os.path.getsize(path_file)
|
||||
if file_size < MIN_FILE_SIZE:
|
||||
os.remove(path_file)
|
||||
dihapus_total += 1
|
||||
stat_count["too_small_file"] += 1
|
||||
log.info(f" [HAPUS-small] {path_file} ({file_size} bytes)")
|
||||
continue
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
try:
|
||||
img = Image.open(path_file)
|
||||
img.verify()
|
||||
except (UnidentifiedImageError, OSError, SyntaxError):
|
||||
try:
|
||||
os.remove(path_file)
|
||||
dihapus_total += 1
|
||||
stat_count["cannot_open"] += 1
|
||||
log.info(f" [HAPUS-corrupt] {path_file}")
|
||||
except OSError:
|
||||
pass
|
||||
continue
|
||||
|
||||
try:
|
||||
img = Image.open(path_file)
|
||||
w, h = img.size
|
||||
if w < MIN_WIDTH or h < MIN_HEIGHT:
|
||||
os.remove(path_file)
|
||||
dihapus_total += 1
|
||||
stat_count["too_small_dims"] += 1
|
||||
log.info(f" [HAPUS-dims] {path_file} ({w}x{h})")
|
||||
continue
|
||||
|
||||
aspect = w / max(h, 1)
|
||||
if aspect > MAX_ASPECT_RATIO or aspect < (1.0 / MAX_ASPECT_RATIO):
|
||||
os.remove(path_file)
|
||||
dihapus_total += 1
|
||||
stat_count["extreme_aspect"] += 1
|
||||
log.info(f" [HAPUS-aspect] {path_file} ({w}x{h}, ar={aspect:.2f})")
|
||||
continue
|
||||
|
||||
if img.mode not in ('RGB', 'RGBA'):
|
||||
img = img.convert('RGB')
|
||||
img.save(path_file)
|
||||
except Exception:
|
||||
try:
|
||||
os.remove(path_file)
|
||||
dihapus_total += 1
|
||||
stat_count["cannot_open"] += 1
|
||||
log.info(f" [HAPUS-exc] {path_file}")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
log.info(f" [OK] Total dihapus: {dihapus_total}")
|
||||
for reason, count in stat_count.items():
|
||||
if count > 0:
|
||||
log.info(f" {reason}: {count}")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 7: PERCEPTUAL HASH DEDUPLICATION
|
||||
# ==========================================
|
||||
def deteksi_duplikat_perceptual():
|
||||
try:
|
||||
import imagehash
|
||||
except ImportError:
|
||||
log.info("--- TAHAP 7: Deteksi Duplikat (SKIP - imagehash not installed) ---")
|
||||
log.info(" Install with: pip install imagehash")
|
||||
return
|
||||
|
||||
log.info("--- TAHAP 7: Deteksi Duplikat Perceptual Hash ---")
|
||||
THRESHOLD = 5
|
||||
if not os.path.exists(TARGET_DIR):
|
||||
log.warning(" [SKIP] Dataset folder tidak ditemukan.")
|
||||
return
|
||||
|
||||
seen_hashes = {}
|
||||
removed = 0
|
||||
scanned = 0
|
||||
|
||||
for class_name in sorted(os.listdir(TARGET_DIR)):
|
||||
class_path = os.path.join(TARGET_DIR, class_name)
|
||||
if not os.path.isdir(class_path):
|
||||
continue
|
||||
for nama_file in sorted(os.listdir(class_path)):
|
||||
path_file = os.path.join(class_path, nama_file)
|
||||
if not os.path.isfile(path_file):
|
||||
continue
|
||||
scanned += 1
|
||||
try:
|
||||
img = Image.open(path_file).convert('RGB')
|
||||
ahash = imagehash.average_hash(img)
|
||||
ahash_hex = str(ahash)
|
||||
is_dup = False
|
||||
for seen_hex, (seen_path, seen_class) in seen_hashes.items():
|
||||
seen_hash = imagehash.hex_to_hash(seen_hex)
|
||||
if ahash - seen_hash <= THRESHOLD:
|
||||
try:
|
||||
os.remove(path_file)
|
||||
removed += 1
|
||||
is_dup = True
|
||||
log.info(f" [DUP] {path_file} ~ {seen_path} (dist={ahash - seen_hash})")
|
||||
break
|
||||
except OSError:
|
||||
pass
|
||||
if not is_dup:
|
||||
seen_hashes[ahash_hex] = (path_file, class_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
log.info(f" [OK] Scanned {scanned}, removed {removed} near-duplicates (threshold={THRESHOLD}).")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TAHAP 8: BUNGKUS KE ZIP
|
||||
# ==========================================
|
||||
def zip_dataset():
|
||||
print("--- TAHAP 5: Mengompresi Folder Dataset ---")
|
||||
log.info("--- TAHAP 8: Mengompresi Folder Dataset ---")
|
||||
if os.path.exists(TARGET_DIR):
|
||||
print(f" Membuat file {TARGET_DIR}.zip, mohon tunggu sebentar...")
|
||||
# shutil.make_archive(nama_output_tanpa_ext, format, folder_yang_dizip)
|
||||
log.info(f" Membuat file {TARGET_DIR}.zip, mohon tunggu sebentar...")
|
||||
shutil.make_archive(TARGET_DIR, 'zip', TARGET_DIR)
|
||||
print(f" [OK] Berhasil! File '{TARGET_DIR}.zip' sudah siap.\n")
|
||||
log.info(f" [OK] Berhasil! File '{TARGET_DIR}.zip' sudah siap.")
|
||||
else:
|
||||
print(f" [ERROR] Folder '{TARGET_DIR}' tidak ditemukan, proses zip dibatalkan.\n")
|
||||
log.error(f" [ERROR] Folder '{TARGET_DIR}' tidak ditemukan, proses zip dibatalkan.")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# MAIN EXECUTION
|
||||
# ==========================================
|
||||
if __name__ == "__main__":
|
||||
print("=== MEMULAI PREPROCESSING DATASET ===\n")
|
||||
log.info("=== MEMULAI PREPROCESSING DATASET ===")
|
||||
|
||||
# Auto-download if ZIPs missing
|
||||
download_dataset("dataset_1.zip", download_dataset_1)
|
||||
download_dataset("dataset_2.zip", download_dataset_2)
|
||||
download_dataset("dataset_3.zip", download_dataset_3)
|
||||
|
||||
ekstrak_semua_zip()
|
||||
gabungkan_dataset_1_dan_2()
|
||||
gabungkan_dataset_3()
|
||||
bersihkan_dataset()
|
||||
hapus_augmented_duplicates()
|
||||
validasi_gambar()
|
||||
deteksi_duplikat_perceptual()
|
||||
zip_dataset()
|
||||
print("=== PREPROCESSING SELESAI ===")
|
||||
print(f"Dataset akhir Anda kini siap digunakan di dalam folder '{TARGET_DIR}' dan '{TARGET_DIR}.zip'.")
|
||||
|
||||
log.info("=== PREPROCESSING SELESAI ===")
|
||||
|
||||
if os.path.exists(TARGET_DIR):
|
||||
total = 0
|
||||
for class_name in sorted(os.listdir(TARGET_DIR)):
|
||||
class_path = os.path.join(TARGET_DIR, class_name)
|
||||
if os.path.isdir(class_path):
|
||||
count = len([f for f in os.listdir(class_path) if os.path.isfile(os.path.join(class_path, f))])
|
||||
total += count
|
||||
log.info(f" {class_name}: {count} images")
|
||||
log.info(f" TOTAL: {total} images")
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
tensorflow==2.19.0 # CPU + Colab; for local GPU, install tensorflow[and-cuda]
|
||||
tensorflowjs==4.22.0
|
||||
gdown # download dataset from Google Drive
|
||||
kagglehub # download dataset from Kaggle
|
||||
numpy
|
||||
matplotlib
|
||||
seaborn
|
||||
|
||||
+138
-34
@@ -1,70 +1,174 @@
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
import traceback
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras import layers, models
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
def build_clean_model(num_classes, img_size=(224, 224)):
|
||||
|
||||
def cbam_block(x, ratio=8, name="cbam"):
|
||||
"""Convolutional Block Attention Module — lightweight foreground attention."""
|
||||
channels = x.shape[-1]
|
||||
|
||||
# Channel attention
|
||||
avg_pool = layers.GlobalAveragePooling2D()(x)
|
||||
max_pool = layers.GlobalMaxPooling2D()(x)
|
||||
ca = layers.Dense(channels // ratio, activation="swish", name=f"{name}_ca1")(avg_pool)
|
||||
ca = layers.Dense(channels, activation="sigmoid", name=f"{name}_ca2")(ca)
|
||||
ca2 = layers.Dense(channels // ratio, activation="swish", name=f"{name}_ca3")(max_pool)
|
||||
ca2 = layers.Dense(channels, activation="sigmoid", name=f"{name}_ca4")(ca2)
|
||||
ca_out = layers.Add(name=f"{name}_ca_add")([ca, ca2])
|
||||
ca_out = layers.Reshape((1, 1, channels), name=f"{name}_ca_reshape")(ca_out)
|
||||
x = layers.Multiply(name=f"{name}_ca_mul")([x, ca_out])
|
||||
|
||||
# Spatial attention
|
||||
from keras import ops
|
||||
avg_sp = ops.mean(x, axis=-1, keepdims=True)
|
||||
max_sp = ops.max(x, axis=-1, keepdims=True)
|
||||
sp = layers.Concatenate(name=f"{name}_sa_cat")([avg_sp, max_sp])
|
||||
sp = layers.Conv2D(1, 7, padding="same", activation="sigmoid", name=f"{name}_sa_conv")(sp)
|
||||
x = layers.Multiply(name=f"{name}_sa_mul")([x, sp])
|
||||
return x
|
||||
|
||||
|
||||
def build_clean_model(num_classes, target_size=(224, 224)):
|
||||
"""Build the production architecture: CBAM + lightweight head, outputting raw logits.
|
||||
|
||||
Mirrors the notebook's build_model() exactly so set_weights() maps correctly.
|
||||
"""
|
||||
base_model = tf.keras.applications.EfficientNetV2B0(
|
||||
input_shape=img_size + (3,),
|
||||
input_shape=target_size + (3,),
|
||||
include_top=False,
|
||||
weights=None,
|
||||
)
|
||||
inputs = tf.keras.Input(shape=img_size + (3,))
|
||||
x = base_model(inputs, training=False)
|
||||
x = layers.Conv2D(512, (3, 3), padding='same', activation='swish')(x)
|
||||
x = layers.BatchNormalization()(x)
|
||||
x = layers.MaxPooling2D((2, 2))(x)
|
||||
x = layers.Dropout(0.2)(x)
|
||||
x = layers.Conv2D(256, (3, 3), padding='same', activation='swish')(x)
|
||||
x = layers.BatchNormalization()(x)
|
||||
x = layers.GlobalAveragePooling2D()(x)
|
||||
x = layers.Dropout(0.3)(x)
|
||||
x = layers.Dense(1024, activation='swish')(x)
|
||||
x = layers.BatchNormalization()(x)
|
||||
x = layers.Dropout(0.4)(x)
|
||||
outputs = layers.Dense(num_classes, activation='softmax', dtype='float32')(x)
|
||||
base_model.trainable = False
|
||||
|
||||
inputs = tf.keras.Input(shape=(None, None, 3), name="input")
|
||||
x = layers.Resizing(target_size[0], target_size[1], interpolation="bilinear",
|
||||
name="resize_input")(inputs)
|
||||
x = layers.GaussianNoise(0.05, name="gauss_noise")(x)
|
||||
x = base_model(x, training=False)
|
||||
# CBAM attention — focus on leaf regions, ignore background
|
||||
x = cbam_block(x, ratio=8, name="cbam")
|
||||
x = layers.GlobalAveragePooling2D(name="gap")(x)
|
||||
x = layers.Dropout(0.3, name="drop_gap")(x)
|
||||
x = layers.Dense(512, activation="swish", name="dense_head")(x)
|
||||
x = layers.BatchNormalization(name="bn_head")(x)
|
||||
x = layers.Dropout(0.4, name="drop_head")(x)
|
||||
# Raw logits (no softmax) — temperature scaling applied at inference
|
||||
outputs = layers.Dense(num_classes, activation="linear", dtype="float32", name="logits")(x)
|
||||
return models.Model(inputs, outputs)
|
||||
|
||||
logging.info("=== EXPORT STARTED ===")
|
||||
|
||||
def _read_labels_for_classes():
|
||||
"""Detect number of classes from model/labels.json or best_model/calibration.json."""
|
||||
for path in ["model/labels.json", "best_model/calibration.json"]:
|
||||
if os.path.exists(path):
|
||||
with open(path) as f:
|
||||
meta = json.load(f)
|
||||
labels = meta.get("labels")
|
||||
if labels:
|
||||
return len(labels)
|
||||
log.warning("Could not detect num_classes from metadata; defaulting to 4.")
|
||||
return 4
|
||||
|
||||
|
||||
log.info("=== EXPORT STARTED (v3.0) ===")
|
||||
try:
|
||||
MODEL_KERAS_PATH = "best_model/best_model.keras"
|
||||
CKPT_DIR = "best_model"
|
||||
WEIGHTS_PATH = os.path.join(CKPT_DIR, "model.weights.h5")
|
||||
MODEL_KERAS_PATH = os.path.join(CKPT_DIR, "best_model.keras")
|
||||
OUTPUT_DIR = "model"
|
||||
saved_model_dir = os.path.join(OUTPUT_DIR, "saved_model")
|
||||
tflite_path = os.path.join(OUTPUT_DIR, "model.tflite")
|
||||
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
if not os.path.exists(MODEL_KERAS_PATH):
|
||||
raise FileNotFoundError(f"Model file not found at {MODEL_KERAS_PATH}")
|
||||
# Use weights H5 as primary source (portable, no Lambda serialization issues)
|
||||
if not os.path.exists(WEIGHTS_PATH):
|
||||
raise FileNotFoundError(
|
||||
f"Weights file not found at {WEIGHTS_PATH}. "
|
||||
"Run the notebook Cell 31 (Save Model) first to generate it."
|
||||
)
|
||||
|
||||
logging.info(f"Loading trained weights from {MODEL_KERAS_PATH}...")
|
||||
original_model = tf.keras.models.load_model(MODEL_KERAS_PATH, compile=False)
|
||||
|
||||
logging.info("Building clean architecture...")
|
||||
clean_model = build_clean_model(num_classes=original_model.output_shape[-1])
|
||||
clean_model.set_weights(original_model.get_weights())
|
||||
logging.info("Weights cloned successfully.")
|
||||
log.info(f"Detecting model configuration...")
|
||||
num_classes = _read_labels_for_classes()
|
||||
|
||||
logging.info(f"Exporting to SavedModel format at: {saved_model_dir}...")
|
||||
log.info(f"Building export architecture ({num_classes} classes)...")
|
||||
clean_model = build_clean_model(num_classes=num_classes)
|
||||
|
||||
log.info(f"Loading trained weights from {WEIGHTS_PATH}...")
|
||||
clean_model.load_weights(WEIGHTS_PATH)
|
||||
log.info("Weights loaded successfully.")
|
||||
|
||||
|
||||
# ─── Export SavedModel (raw logits) ───
|
||||
log.info(f"Exporting SavedModel (raw logits) at: {saved_model_dir}...")
|
||||
tf.saved_model.save(clean_model, saved_model_dir)
|
||||
logging.info("SavedModel export completed successfully.")
|
||||
log.info("SavedModel export completed successfully.")
|
||||
|
||||
# ─── Export TFLite (INT8 quantization) ───
|
||||
log.info(f"Converting to TFLite INT8 at: {tflite_path}...")
|
||||
|
||||
# Representative dataset for INT8 quantization
|
||||
def representative_dataset():
|
||||
val_dir = "dataset_split/val"
|
||||
if not os.path.exists(val_dir):
|
||||
log.warning("Validation dir not found; skipping representative dataset.")
|
||||
return
|
||||
ds = tf.keras.utils.image_dataset_from_directory(
|
||||
val_dir, shuffle=True, batch_size=1, image_size=(224, 224)
|
||||
)
|
||||
for images, _ in ds.take(200):
|
||||
yield [tf.cast(images, tf.float32)]
|
||||
|
||||
logging.info(f"Converting to TFLite format at: {tflite_path}...")
|
||||
converter = tf.lite.TFLiteConverter.from_keras_model(clean_model)
|
||||
converter.target_spec.supported_ops = [
|
||||
tf.lite.OpsSet.TFLITE_BUILTINS,
|
||||
tf.lite.OpsSet.SELECT_TF_OPS,
|
||||
]
|
||||
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
||||
if os.path.exists("dataset_split/val"):
|
||||
converter.representative_dataset = representative_dataset
|
||||
|
||||
tflite_model = converter.convert()
|
||||
with open(tflite_path, "wb") as f:
|
||||
f.write(tflite_model)
|
||||
logging.info("TFLite conversion completed successfully.")
|
||||
logging.info("=== EXPORT COMPLETED ===")
|
||||
log.info("TFLite conversion completed successfully.")
|
||||
|
||||
# ─── Export model metadata ───
|
||||
# Load labels and calibration from training output
|
||||
labels_path = os.path.join(OUTPUT_DIR, "labels.json")
|
||||
cal_path = os.path.join("best_model", "calibration.json")
|
||||
|
||||
labels_meta = {"labels": None, "temperature": 1.0, "conf_threshold_high": 0.70,
|
||||
"conf_threshold_low": 0.45}
|
||||
|
||||
if os.path.exists(labels_path):
|
||||
with open(labels_path) as f:
|
||||
labels_meta.update(json.load(f))
|
||||
|
||||
if os.path.exists(cal_path):
|
||||
with open(cal_path) as f:
|
||||
cal = json.load(f)
|
||||
labels_meta["temperature"] = cal.get("temperature", 1.0)
|
||||
|
||||
labels_meta["version"] = "3.0"
|
||||
labels_meta["architecture"] = "EfficientNetV2B0 + CBAM + Dense(512)"
|
||||
labels_meta["output_type"] = "logits"
|
||||
labels_meta["input_range"] = [0, 255]
|
||||
labels_meta["input_size"] = [224, 224]
|
||||
labels_meta["preprocessing"] = "resize_bilinear_224x224_no_normalization"
|
||||
|
||||
with open(labels_path, "w") as f:
|
||||
json.dump(labels_meta, f, indent=2)
|
||||
log.info(f"Labels + calibration metadata saved to {labels_path}")
|
||||
|
||||
log.info("=== EXPORT COMPLETED (v3.0) ===")
|
||||
|
||||
except Exception:
|
||||
logging.error("EXPORT FAILED")
|
||||
logging.error(traceback.format_exc())
|
||||
log.error("EXPORT FAILED")
|
||||
log.error(traceback.format_exc())
|
||||
|
||||
@@ -2,18 +2,24 @@ use anyhow::{Context, Result};
|
||||
use std::env;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub const LABELS: [&str; 4] = ["Bercak Daun", "Daun Sehat", "Karat Daun", "Hawar Daun"];
|
||||
pub const LABELS: [&str; 4] = ["Bercak Daun", "Daun Sehat", "Hawar Daun", "Karat Daun"];
|
||||
pub const SERVICE_NAME: &str = "zeavis-ml-service";
|
||||
pub const SERVICE_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
pub const DEFAULT_INPUT_SIZE: u32 = 224;
|
||||
pub const DEFAULT_MODEL_PATH: &str = "../../Machine_Learning/model/model.onnx";
|
||||
pub const DEFAULT_TEMPERATURE: f32 = 1.0;
|
||||
pub const CONFIDENCE_THRESHOLD_HIGH: f32 = 0.70;
|
||||
pub const CONFIDENCE_THRESHOLD_LOW: f32 = 0.45;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Config {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub model_path: PathBuf,
|
||||
pub input_size: u32,
|
||||
pub temperature: f32,
|
||||
pub conf_threshold_high: f32,
|
||||
pub conf_threshold_low: f32,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
@@ -27,12 +33,18 @@ impl Config {
|
||||
let port = parse_env_u16("ML_SERVICE_PORT", 8000)?;
|
||||
let input_size = parse_env_u32("MODEL_INPUT_SIZE", DEFAULT_INPUT_SIZE)?;
|
||||
let model_path = env::var("MODEL_PATH").unwrap_or_else(|_| DEFAULT_MODEL_PATH.to_string());
|
||||
let temperature = parse_env_f32("MODEL_TEMPERATURE", DEFAULT_TEMPERATURE)?;
|
||||
let conf_threshold_high = parse_env_f32("MODEL_CONF_HIGH", CONFIDENCE_THRESHOLD_HIGH)?;
|
||||
let conf_threshold_low = parse_env_f32("MODEL_CONF_LOW", CONFIDENCE_THRESHOLD_LOW)?;
|
||||
|
||||
Ok(Self {
|
||||
host,
|
||||
port,
|
||||
model_path: resolve_model_path(base_dir, &model_path),
|
||||
input_size,
|
||||
temperature,
|
||||
conf_threshold_high,
|
||||
conf_threshold_low,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -64,13 +76,22 @@ fn parse_env_u32(name: &str, default: u32) -> Result<u32> {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_env_f32(name: &str, default: f32) -> Result<f32> {
|
||||
match env::var(name) {
|
||||
Ok(value) => value
|
||||
.parse::<f32>()
|
||||
.with_context(|| format!("{name} must be a valid f32")),
|
||||
Err(_) => Ok(default),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn labels_match_training_class_order_with_display_names() {
|
||||
assert_eq!(LABELS, ["Bercak Daun", "Daun Sehat", "Karat Daun", "Hawar Daun"]);
|
||||
assert_eq!(LABELS, ["Bercak Daun", "Daun Sehat", "Hawar Daun", "Karat Daun"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -23,13 +23,20 @@ async fn main() -> Result<()> {
|
||||
// Load configuration from environment
|
||||
let config = Config::from_env()?;
|
||||
|
||||
// Create ModelService and wrap in Arc
|
||||
let model = Arc::new(ModelService::new(&config.model_path, config.input_size));
|
||||
|
||||
// Log model status
|
||||
// Create ModelService with calibration and wrap in Arc
|
||||
let model = Arc::new(ModelService::with_calibration(
|
||||
&config.model_path,
|
||||
config.input_size,
|
||||
config.temperature,
|
||||
config.conf_threshold_high,
|
||||
config.conf_threshold_low,
|
||||
));
|
||||
tracing::info!(
|
||||
model_loaded = model.is_loaded(),
|
||||
model_path = ?config.model_path,
|
||||
temperature = config.temperature,
|
||||
conf_high = config.conf_threshold_high,
|
||||
conf_low = config.conf_threshold_low,
|
||||
"Model service initialized"
|
||||
);
|
||||
|
||||
|
||||
+125
-117
@@ -1,4 +1,4 @@
|
||||
use crate::config::LABELS;
|
||||
use crate::config::{CONFIDENCE_THRESHOLD_HIGH, CONFIDENCE_THRESHOLD_LOW, DEFAULT_TEMPERATURE, LABELS};
|
||||
use crate::error::ServiceError;
|
||||
use ndarray::Array4;
|
||||
use ort::{session::Session, value::TensorRef};
|
||||
@@ -7,32 +7,40 @@ use std::collections::BTreeMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Mutex;
|
||||
|
||||
/// Prediction result containing the top label, confidence, and all probabilities.
|
||||
/// Prediction result with temperature-calibrated probabilities and status.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct Prediction {
|
||||
pub status: String, // "confident", "uncertain", "rejected"
|
||||
pub label: String,
|
||||
pub confidence: f32,
|
||||
pub probabilities: BTreeMap<String, f32>,
|
||||
}
|
||||
|
||||
/// Service for running ONNX model inference.
|
||||
/// Service for running ONNX model inference with temperature-scaled calibration.
|
||||
///
|
||||
/// Stores the model path, input size, and an optional thread-safe ONNX session.
|
||||
/// If the model fails to load, the session remains None and predictions will fail.
|
||||
/// The session is wrapped in a Mutex to ensure thread-safe access from concurrent Axum requests.
|
||||
/// The ONNX model outputs raw logits. Temperature scaling + softmax is applied
|
||||
/// in predict() to produce calibrated probabilities and a decision status:
|
||||
/// - confident: max_prob >= conf_threshold_high
|
||||
/// - uncertain: conf_threshold_low <= max_prob < conf_threshold_high
|
||||
/// - rejected: max_prob < conf_threshold_low
|
||||
pub struct ModelService {
|
||||
model_path: std::path::PathBuf,
|
||||
input_size: u32,
|
||||
temperature: f32,
|
||||
conf_threshold_high: f32,
|
||||
conf_threshold_low: f32,
|
||||
session: Option<Mutex<Session>>,
|
||||
}
|
||||
|
||||
impl ModelService {
|
||||
/// Creates a new ModelService, attempting to load the ONNX model from the given path.
|
||||
///
|
||||
/// If the model file does not exist or fails to load, the session is stored as None.
|
||||
/// This allows the service to report unloaded state via health checks.
|
||||
/// The session is wrapped in a Mutex for thread-safe concurrent access.
|
||||
pub fn new(model_path: &Path, input_size: u32) -> Self {
|
||||
Self::with_calibration(model_path, input_size, DEFAULT_TEMPERATURE,
|
||||
CONFIDENCE_THRESHOLD_HIGH, CONFIDENCE_THRESHOLD_LOW)
|
||||
}
|
||||
|
||||
/// Creates a new ModelService with temperature scaling and confidence thresholds.
|
||||
pub fn with_calibration(model_path: &Path, input_size: u32,
|
||||
temperature: f32, conf_high: f32, conf_low: f32) -> Self {
|
||||
let session = Session::builder()
|
||||
.ok()
|
||||
.and_then(|mut builder| builder.commit_from_file(model_path).ok())
|
||||
@@ -41,88 +49,101 @@ impl ModelService {
|
||||
Self {
|
||||
model_path: model_path.to_path_buf(),
|
||||
input_size,
|
||||
temperature,
|
||||
conf_threshold_high: conf_high,
|
||||
conf_threshold_low: conf_low,
|
||||
session,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if the model is loaded and ready for inference.
|
||||
pub fn is_loaded(&self) -> bool {
|
||||
self.session.is_some()
|
||||
}
|
||||
|
||||
/// Returns the path to the model file.
|
||||
pub fn model_path(&self) -> &Path {
|
||||
&self.model_path
|
||||
}
|
||||
|
||||
/// Returns the input size (width/height) for the model.
|
||||
pub fn input_size(&self) -> u32 {
|
||||
self.input_size
|
||||
}
|
||||
|
||||
/// Runs inference on the given input array.
|
||||
pub fn temperature(&self) -> f32 {
|
||||
self.temperature
|
||||
}
|
||||
|
||||
/// Runs inference and returns temperature-calibrated Prediction.
|
||||
///
|
||||
/// Returns ModelUnavailable if the model is not loaded.
|
||||
/// Returns PredictionFailed if inference fails, lock is poisoned, or output format is invalid.
|
||||
/// The ONNX model outputs raw logits (no softmax). Temperature scaling
|
||||
/// is applied: probs = softmax(logits / T).
|
||||
pub fn predict(&self, input: Array4<f32>) -> Result<Prediction, ServiceError> {
|
||||
let session = self
|
||||
.session
|
||||
.as_ref()
|
||||
let session = self.session.as_ref()
|
||||
.ok_or_else(|| ServiceError::ModelUnavailable("Model is not loaded".to_string()))?;
|
||||
|
||||
// Lock the session for thread-safe access
|
||||
let mut session_guard = session
|
||||
.lock()
|
||||
let mut session_guard = session.lock()
|
||||
.map_err(|_| ServiceError::PredictionFailed("Prediction failed".to_string()))?;
|
||||
|
||||
let input = TensorRef::from_array_view(&input)
|
||||
.map_err(|_| ServiceError::PredictionFailed("Prediction failed".to_string()))?;
|
||||
let outputs = session_guard
|
||||
.run(ort::inputs![input])
|
||||
let outputs = session_guard.run(ort::inputs![input])
|
||||
.map_err(|_| ServiceError::PredictionFailed("Prediction failed".to_string()))?;
|
||||
|
||||
let output_tensor = outputs[0]
|
||||
.try_extract_tensor::<f32>()
|
||||
let output_tensor = outputs[0].try_extract_tensor::<f32>()
|
||||
.map_err(|_| ServiceError::PredictionFailed("Prediction failed".to_string()))?;
|
||||
|
||||
let probabilities: Vec<f32> = output_tensor.1.iter().copied().collect();
|
||||
let logits: Vec<f32> = output_tensor.1.iter().copied().collect();
|
||||
|
||||
Self::prediction_from_probabilities(&probabilities)
|
||||
Self::calibrate_prediction(&logits, self.temperature,
|
||||
self.conf_threshold_high, self.conf_threshold_low)
|
||||
}
|
||||
|
||||
/// Maps a probability vector to a Prediction with label and all probabilities.
|
||||
///
|
||||
/// Expects a vector of length 4 (one per label in LABELS).
|
||||
/// Rejects non-finite values (NaN, +inf, -inf) to prevent invalid predictions.
|
||||
/// Returns PredictionFailed if the length is incorrect or any value is non-finite.
|
||||
pub fn prediction_from_probabilities(probs: &[f32]) -> Result<Prediction, ServiceError> {
|
||||
if probs.len() != LABELS.len() {
|
||||
/// Applies temperature scaling + softmax to raw logits, determines status.
|
||||
fn calibrate_prediction(logits: &[f32], temperature: f32,
|
||||
conf_high: f32, conf_low: f32) -> Result<Prediction, ServiceError> {
|
||||
if logits.len() != LABELS.len() {
|
||||
return Err(ServiceError::PredictionFailed("Prediction failed".to_string()));
|
||||
}
|
||||
|
||||
// Reject non-finite values (NaN, +inf, -inf)
|
||||
if probs.iter().any(|p| !p.is_finite()) {
|
||||
// Reject non-finite values
|
||||
if logits.iter().any(|p| !p.is_finite()) {
|
||||
return Err(ServiceError::PredictionFailed("Prediction failed".to_string()));
|
||||
}
|
||||
|
||||
// Find the index with the highest probability
|
||||
let top_idx = probs
|
||||
.iter()
|
||||
.enumerate()
|
||||
// Temperature scaling: divide by T
|
||||
let T = if temperature > 0.0 { temperature } else { 1.0 };
|
||||
let scaled: Vec<f32> = logits.iter().map(|l| l / T).collect();
|
||||
|
||||
// Numerically stable softmax: shift by max to avoid overflow
|
||||
let max_logit = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exp_vals: Vec<f32> = scaled.iter().map(|l| (l - max_logit).exp()).collect();
|
||||
let sum: f32 = exp_vals.iter().sum();
|
||||
let probs: Vec<f32> = exp_vals.iter().map(|e| e / sum).collect();
|
||||
|
||||
// Find top probability
|
||||
let top_idx = probs.iter().enumerate()
|
||||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.map(|(idx, _)| idx)
|
||||
.unwrap_or(0);
|
||||
|
||||
let top_label = LABELS[top_idx].to_string();
|
||||
let confidence = probs[top_idx];
|
||||
let top_label = LABELS[top_idx].to_string();
|
||||
|
||||
// Determine status
|
||||
let status = if confidence >= conf_high {
|
||||
"confident"
|
||||
} else if confidence >= conf_low {
|
||||
"uncertain"
|
||||
} else {
|
||||
"rejected"
|
||||
};
|
||||
|
||||
// Build probabilities map
|
||||
let mut probabilities = BTreeMap::new();
|
||||
for (i, &prob) in probs.iter().enumerate() {
|
||||
probabilities.insert(LABELS[i].to_string(), prob);
|
||||
}
|
||||
|
||||
Ok(Prediction {
|
||||
status: status.to_string(),
|
||||
label: top_label,
|
||||
confidence,
|
||||
probabilities,
|
||||
@@ -134,105 +155,92 @@ impl ModelService {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn prediction_mapping_selects_top_label_and_all_probabilities() {
|
||||
let probs = [0.1, 0.2, 0.6, 0.1];
|
||||
let result = ModelService::prediction_from_probabilities(&probs);
|
||||
const T: f32 = 1.0;
|
||||
const HIGH: f32 = 0.70;
|
||||
const LOW: f32 = 0.45;
|
||||
|
||||
#[test]
|
||||
fn calibrate_probs_selects_top_label() {
|
||||
// [1, 2, 3, 0.5] → softmax ≈ [0.086, 0.235, 0.638, 0.040]
|
||||
// 0.638 < 0.70 → "uncertain". Use larger gap for "confident".
|
||||
let logits = [0.0, 0.0, 10.0, 0.0]; // softmax ≈ [0, 0, ~1, 0]
|
||||
let result = ModelService::calibrate_prediction(&logits, T, HIGH, LOW);
|
||||
assert!(result.is_ok());
|
||||
let prediction = result.unwrap();
|
||||
|
||||
// Top label should be Karat Daun (index 2 with 0.6 probability)
|
||||
assert_eq!(prediction.label, "Karat Daun");
|
||||
assert_eq!(prediction.confidence, 0.6);
|
||||
|
||||
// All probabilities should be present
|
||||
assert_eq!(prediction.probabilities.len(), 4);
|
||||
assert_eq!(prediction.probabilities.get("Bercak Daun"), Some(&0.1));
|
||||
assert_eq!(prediction.probabilities.get("Daun Sehat"), Some(&0.2));
|
||||
assert_eq!(prediction.probabilities.get("Karat Daun"), Some(&0.6));
|
||||
assert_eq!(prediction.probabilities.get("Hawar Daun"), Some(&0.1));
|
||||
let p = result.unwrap();
|
||||
// Index 2 = Hawar Daun (highest logit)
|
||||
assert_eq!(p.label, "Hawar Daun");
|
||||
assert!(p.confidence >= 0.999);
|
||||
assert_eq!(p.status, "confident");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prediction_mapping_rejects_wrong_output_length() {
|
||||
let probs = [0.25, 0.25, 0.25]; // Only 3 values instead of 4
|
||||
let result = ModelService::prediction_from_probabilities(&probs);
|
||||
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
ServiceError::PredictionFailed(msg) => {
|
||||
assert_eq!(msg, "Prediction failed");
|
||||
}
|
||||
_ => panic!("expected PredictionFailed error"),
|
||||
}
|
||||
fn calibrate_probs_uncertain_when_borderline() {
|
||||
let logits = [0.0, 0.4, 0.0, 0.0]; // softmax with low max
|
||||
let result = ModelService::calibrate_prediction(&logits, 2.0, HIGH, LOW);
|
||||
assert!(result.is_ok());
|
||||
let p = result.unwrap();
|
||||
// T=2.0 flattens further — likely uncertain or rejected
|
||||
assert!(p.status == "uncertain" || p.status == "rejected");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prediction_mapping_rejects_nan_values() {
|
||||
let probs = [0.1, f32::NAN, 0.6, 0.1];
|
||||
let result = ModelService::prediction_from_probabilities(&probs);
|
||||
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
ServiceError::PredictionFailed(msg) => {
|
||||
assert_eq!(msg, "Prediction failed");
|
||||
}
|
||||
_ => panic!("expected PredictionFailed error"),
|
||||
}
|
||||
fn calibrate_probs_rejects_low_confidence() {
|
||||
let logits = [0.01, 0.01, 0.01, 0.02];
|
||||
let result = ModelService::calibrate_prediction(&logits, 10.0, HIGH, LOW);
|
||||
assert!(result.is_ok());
|
||||
let p = result.unwrap();
|
||||
assert_eq!(p.status, "rejected");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prediction_mapping_rejects_positive_infinity() {
|
||||
let probs = [0.1, 0.2, f32::INFINITY, 0.1];
|
||||
let result = ModelService::prediction_from_probabilities(&probs);
|
||||
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
ServiceError::PredictionFailed(msg) => {
|
||||
assert_eq!(msg, "Prediction failed");
|
||||
}
|
||||
_ => panic!("expected PredictionFailed error"),
|
||||
}
|
||||
fn calibrate_probs_all_probabilities_present() {
|
||||
let logits = [1.0, 2.0, 3.0, 4.0];
|
||||
let result = ModelService::calibrate_prediction(&logits, T, HIGH, LOW);
|
||||
assert!(result.is_ok());
|
||||
let p = result.unwrap();
|
||||
assert_eq!(p.probabilities.len(), 4);
|
||||
assert!(p.probabilities.contains_key("Bercak Daun"));
|
||||
assert!(p.probabilities.contains_key("Daun Sehat"));
|
||||
assert!(p.probabilities.contains_key("Hawar Daun"));
|
||||
assert!(p.probabilities.contains_key("Karat Daun"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prediction_mapping_rejects_negative_infinity() {
|
||||
let probs = [0.1, 0.2, 0.6, f32::NEG_INFINITY];
|
||||
let result = ModelService::prediction_from_probabilities(&probs);
|
||||
|
||||
fn calibrate_probs_rejects_wrong_length() {
|
||||
let logits = [0.25, 0.25, 0.25];
|
||||
let result = ModelService::calibrate_prediction(&logits, T, HIGH, LOW);
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
ServiceError::PredictionFailed(msg) => {
|
||||
assert_eq!(msg, "Prediction failed");
|
||||
}
|
||||
_ => panic!("expected PredictionFailed error"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn calibrate_probs_rejects_nan() {
|
||||
let logits = [0.1, f32::NAN, 0.6, 0.1];
|
||||
let result = ModelService::calibrate_prediction(&logits, T, HIGH, LOW);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn temperature_one_gives_same_ranking() {
|
||||
let logits = [0.0, 1.0, 2.0, 3.0];
|
||||
let r1 = ModelService::calibrate_prediction(&logits, 1.0, 0.0, 0.0).unwrap();
|
||||
let r2 = ModelService::calibrate_prediction(&logits, 2.0, 0.0, 0.0).unwrap();
|
||||
assert_eq!(r1.label, r2.label);
|
||||
assert!(r1.confidence > r2.confidence); // T=2 flattens
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_model_file_creates_unloaded_service() {
|
||||
let model_path = Path::new("/nonexistent/model.onnx");
|
||||
let service = ModelService::new(model_path, 224);
|
||||
|
||||
let service = ModelService::new(Path::new("/nonexistent/model.onnx"), 224);
|
||||
assert!(!service.is_loaded());
|
||||
assert_eq!(service.model_path(), model_path);
|
||||
assert_eq!(service.input_size(), 224);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unloaded_service_returns_model_unavailable() {
|
||||
let model_path = Path::new("/nonexistent/model.onnx");
|
||||
let service = ModelService::new(model_path, 224);
|
||||
|
||||
let dummy_input = ndarray::Array4::zeros((1, 224, 224, 3));
|
||||
let result = service.predict(dummy_input);
|
||||
|
||||
let service = ModelService::new(Path::new("/nonexistent/model.onnx"), 224);
|
||||
let result = service.predict(ndarray::Array4::zeros((1, 224, 224, 3)));
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
ServiceError::ModelUnavailable(msg) => {
|
||||
assert_eq!(msg, "Model is not loaded");
|
||||
}
|
||||
ServiceError::ModelUnavailable(msg) => assert_eq!(msg, "Model is not loaded"),
|
||||
_ => panic!("expected ModelUnavailable error"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -30,6 +30,7 @@ pub struct MetadataResponse {
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct PredictionResponse {
|
||||
pub status: String,
|
||||
pub label: String,
|
||||
pub confidence: f32,
|
||||
pub probabilities: std::collections::BTreeMap<String, f32>,
|
||||
@@ -60,6 +61,7 @@ pub fn metadata_response(model_path: String, model_loaded: bool, input_size: u32
|
||||
|
||||
pub fn prediction_response(prediction: Prediction) -> PredictionResponse {
|
||||
PredictionResponse {
|
||||
status: prediction.status,
|
||||
label: prediction.label,
|
||||
confidence: prediction.confidence,
|
||||
probabilities: prediction.probabilities,
|
||||
@@ -130,7 +132,7 @@ pub async fn predict(
|
||||
};
|
||||
|
||||
// Preprocess the image
|
||||
let preprocess_start = std::time::Instant::now();
|
||||
let _preprocess_start = std::time::Instant::now();
|
||||
let input = preprocess_image(&bytes, state.model.input_size())?;
|
||||
|
||||
// Record image size metric
|
||||
@@ -202,33 +204,35 @@ mod tests {
|
||||
assert_eq!(response.labels.len(), 4);
|
||||
assert_eq!(
|
||||
response.labels,
|
||||
vec!["Bercak Daun", "Daun Sehat", "Karat Daun", "Hawar Daun"]
|
||||
vec!["Bercak Daun", "Daun Sehat", "Hawar Daun", "Karat Daun"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prediction_response_matches_prediction_contract() {
|
||||
let prediction = Prediction {
|
||||
label: "Karat Daun".to_string(),
|
||||
status: "confident".to_string(),
|
||||
label: "Hawar Daun".to_string(),
|
||||
confidence: 0.6,
|
||||
probabilities: {
|
||||
let mut map = std::collections::BTreeMap::new();
|
||||
map.insert("Bercak Daun".to_string(), 0.1);
|
||||
map.insert("Daun Sehat".to_string(), 0.2);
|
||||
map.insert("Karat Daun".to_string(), 0.6);
|
||||
map.insert("Hawar Daun".to_string(), 0.1);
|
||||
map.insert("Hawar Daun".to_string(), 0.6);
|
||||
map.insert("Karat Daun".to_string(), 0.1);
|
||||
map
|
||||
},
|
||||
};
|
||||
|
||||
let response = prediction_response(prediction);
|
||||
|
||||
assert_eq!(response.label, "Karat Daun");
|
||||
assert_eq!(response.status, "confident");
|
||||
assert_eq!(response.label, "Hawar Daun");
|
||||
assert_eq!(response.confidence, 0.6);
|
||||
assert_eq!(response.probabilities.len(), 4);
|
||||
assert_eq!(response.probabilities.get("Bercak Daun"), Some(&0.1));
|
||||
assert_eq!(response.probabilities.get("Daun Sehat"), Some(&0.2));
|
||||
assert_eq!(response.probabilities.get("Karat Daun"), Some(&0.6));
|
||||
assert_eq!(response.probabilities.get("Hawar Daun"), Some(&0.1));
|
||||
assert_eq!(response.probabilities.get("Hawar Daun"), Some(&0.6));
|
||||
assert_eq!(response.probabilities.get("Karat Daun"), Some(&0.1));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user