Mystic Dicky

legacy/test_cam.py

@@ -0,0 +1,202 @@
+import cv2
+import time
+import threading
+from collections import deque
+from ultralytics.models import YOLO
+
+# ========= CẤU HÌNH =========
+CAM_INDEX = 0
+BUFFER_SIZE = 10
+PROCESS_FPS = 2
+DET_THRESH = 0.8
+CLS_THRESH = 0.8
+
+DET_MODEL_PATH = r"F:\Intership\Smoke\test\cig-det-best.pt"
+CLS_MODEL_PATH = r"F:\Intership\Smoke\test\cig-cls-best.pt"
+
+# Nếu classifier có 2 lớp: smoke / no_smoke (hoặc smoking / not_smoking)
+# Bạn đặt tên lớp "hút thuốc" tại đây cho chắc chắn.
+SMOKING_CLASS_NAME = "smoke"  # đổi thành "smoking" nếu model bạn đặt tên như vậy
+
+# ========= BIẾN CHIA SẺ =========
+frame_buffer = deque(maxlen=BUFFER_SIZE)
+buffer_lock = threading.Lock()
+stop_event = threading.Event()
+
+# ========= LOAD MODEL =========
+det_model = YOLO(DET_MODEL_PATH)
+cls_model = YOLO(CLS_MODEL_PATH)
+
+
+def capture_and_display():
+    cap = cv2.VideoCapture(CAM_INDEX)
+    if not cap.isOpened():
+        print("Không mở được camera.")
+        stop_event.set()
+        return
+
+    while not stop_event.is_set():
+        ret, frame = cap.read()
+        if not ret:
+            print("Không đọc được frame từ camera.")
+            break
+
+        with buffer_lock:
+            frame_buffer.append(frame)
+
+        cv2.imshow("Realtime Camera", frame)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            stop_event.set()
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+
+def get_latest_frame():
+    with buffer_lock:
+        if not frame_buffer:
+            return None
+        return frame_buffer[-1].copy()
+
+
+def run_detection(frame):
+    """
+    Return:
+      best_conf: float
+      best_box: (x1,y1,x2,y2) int or None
+    """
+    results = det_model.predict(source=frame, verbose=False)
+    r = results[0]
+
+    if r.boxes is None or len(r.boxes) == 0:
+        return 0.0, None
+
+    confs = r.boxes.conf.cpu().numpy()
+    boxes = r.boxes.xyxy.cpu().numpy()
+
+    best_idx = confs.argmax()
+    best_conf = float(confs[best_idx])
+    x1, y1, x2, y2 = boxes[best_idx]
+    best_box = (int(x1), int(y1), int(x2), int(y2))
+
+    return best_conf, best_box
+
+
+def crop_from_box(frame, box, pad=10):
+    """
+    Crop theo box + padding nhẹ.
+    Nếu box None -> dùng full frame.
+    """
+    if box is None:
+        return frame
+
+    x1, y1, x2, y2 = box
+    h, w = frame.shape[:2]
+
+    x1 = max(0, x1 - pad)
+    y1 = max(0, y1 - pad)
+    x2 = min(w - 1, x2 + pad)
+    y2 = min(h - 1, y2 + pad)
+
+    if x2 <= x1 or y2 <= y1:
+        return frame
+    return frame[y1:y2, x1:x2]
+
+
+def run_classification(img):
+    """
+    Return:
+      cls_name: str
+      cls_conf: float
+    """
+    results = cls_model.predict(source=img, verbose=False)
+    r = results[0]
+
+    probs = r.probs
+    top1 = int(probs.top1)
+    top1conf = float(probs.top1conf)
+
+    names = r.names
+    if isinstance(names, dict):
+        cls_name = names.get(top1, str(top1))
+    else:
+        cls_name = names[top1] if top1 < len(names) else str(top1)
+
+    return cls_name, top1conf
+
+
+def is_smoking_by_class(cls_name, cls_conf):
+    """
+    CHỈ báo hút thuốc khi:
+    - confidence đủ cao
+    - và class đúng là lớp hút thuốc (SMOKING_CLASS_NAME)
+    """
+    if cls_conf < CLS_THRESH:
+        return False
+
+    # bắt buộc phải biết class "hút thuốc" là gì
+    if not SMOKING_CLASS_NAME:
+        # Nếu bạn không set được tên class, KHÔNG được auto-True
+        return False
+
+    return cls_name.strip().lower() == SMOKING_CLASS_NAME.strip().lower()
+
+
+def processing_loop():
+    interval = 1.0 / PROCESS_FPS
+    last_time = 0.0
+
+    while not stop_event.is_set():
+        now = time.time()
+        if now - last_time < interval:
+            time.sleep(0.001)
+            continue
+        last_time = now
+
+        frame = get_latest_frame()
+        if frame is None:
+            continue
+
+        # ===== Step 1: Detection =====
+        det_conf, det_box = run_detection(frame)
+
+        if det_conf >= DET_THRESH:
+            print(f"[ALERT] HÚT THUỐC (det_conf={det_conf:.2f})")
+            continue
+
+        # ===== Step 2: det_conf < 0.8 -> crop -> classification =====
+        crop_img = crop_from_box(
+            frame, det_box
+        )  # nếu có box thì crop, không có thì full frame
+        cls_name, cls_conf = run_classification(crop_img)
+
+        if is_smoking_by_class(cls_name, cls_conf):
+            print(
+                f"[ALERT] CÓ HÚT THUỐC (cls={cls_name}, cls_conf={cls_conf:.2f}, det_conf={det_conf:.2f})"
+            )
+        else:
+            print(
+                f"[OK] KHÔNG HÚT/CHƯA CHẮC (cls={cls_name}, cls_conf={cls_conf:.2f}, det_conf={det_conf:.2f})"
+            )
+
+
+def main():
+    t1 = threading.Thread(target=capture_and_display, daemon=True)
+    t2 = threading.Thread(target=processing_loop, daemon=True)
+
+    t1.start()
+    t2.start()
+
+    try:
+        while not stop_event.is_set():
+            time.sleep(0.1)
+    except KeyboardInterrupt:
+        stop_event.set()
+
+    t1.join()
+    t2.join()
+
+
+if __name__ == "__main__":
+    main()