Spaces:
Sleeping
Sleeping
File size: 2,733 Bytes
248b460 42a5c71 248b460 42a5c71 248b460 42a5c71 248b460 42a5c71 248b460 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | import os
from detectron2.utils.logger import setup_logger
from detectron2.data.datasets import register_coco_instances
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
# Cấu hình log
setup_logger()
# Thiết lập đường dẫn thư mục gốc tương đối
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(CURRENT_DIR)
# Đường dẫn đến tập validation
# VALID_JSON = os.path.join(PROJECT_ROOT, "Datasets", "Data_stage1", "valid", "_annotations.coco.json")
# VALID_DIR = os.path.join(PROJECT_ROOT, "Datasets", "Data_stage1", "valid")
VALID_JSON = os.path.join(PROJECT_ROOT, "Datasets", "Dataset_main", "valid", "_annotations.coco.json")
VALID_DIR = os.path.join(PROJECT_ROOT, "Datasets", "Dataset_main", "valid")
# Nếu tập dataset chưa được đăng ký trong session hiện tại, thực hiện đăng ký
try:
register_coco_instances("tech_draw_valid", {}, VALID_JSON, VALID_DIR)
except Exception as e:
print(f"Dataset đã được đăng ký hoặc gặp lỗi: {e}")
def get_eval_cfg():
cfg = get_cfg()
# Đọc base config giống lúc train
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
# Cấu hình dataset
cfg.DATASETS.TEST = ("tech_draw_valid",)
# Cấu hình kiến trúc giống quá trình training
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.MODEL.MASK_ON = False
cfg.MODEL.DEVICE = "cuda" # hoặc "cpu" nếu chạy test không có GPU
# Load weights model đã train thành công (file model_final.pth nằm trong thư mục output_model)
cfg.MODEL.WEIGHTS = os.path.join(CURRENT_DIR, "output_model_stage2", "model_final.pth")
# Thiết lập ngưỡng confidence score để evaluate (có thể tinh chỉnh, vd: lấy box có tin cậy > 50%)
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.3
return cfg
def main():
cfg = get_eval_cfg()
print("Đang load model...")
predictor = DefaultPredictor(cfg)
print("Bắt đầu Evaluate tính mAP...")
# COCOEvaluator sử dụng format của MS COCO để đo lường mAP
evaluator = COCOEvaluator("tech_draw_valid", output_dir=os.path.join(CURRENT_DIR, "output_model_stage2"))
val_loader = build_detection_test_loader(cfg, "tech_draw_valid")
# Thực thi tính metrics
print(inference_on_dataset(predictor.model, val_loader, evaluator))
if __name__ == "__main__":
main()
|