main.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 

import cv2
import numpy as np
import tensorflow as tf

# 1. Setup Paths
base_path = os.path.dirname(__file__)
labels_path = os.path.join(base_path, "labels.txt")
model_path = os.path.join(base_path, ".") 

# 2. Load Labels & Model
class_names = [line.strip() for line in open(labels_path, "r").readlines()]
model = tf.saved_model.load(model_path)
infer = model.signatures["serving_default"]

camera = cv2.VideoCapture(0)

while True:
    ret, frame = camera.read()
    if not ret: break

    # 3. Pre-process (CRITICAL FOR ACCURACY)
    display_img = cv2.flip(frame, 1)
    h, w, _ = display_img.shape
    
    # Resize and convert to RGB for the AI
    img_rgb = cv2.cvtColor(display_img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, (224, 224), interpolation=cv2.INTER_LINEAR)
    
    # Math: Convert 0-255 pixels to -1 to +1 range
    img_array = np.expand_dims(img_resized, axis=0).astype(np.float32)
    img_array = (img_array / 127.5) - 1.0 

    # 4. Prediction
    output = infer(tf.convert_to_tensor(img_array))
    output_key = list(output.keys())[0]
    predictions = output[output_key].numpy()[0]
    
    index = np.argmax(predictions)
    probability = predictions[index]
    
    # Identify the detected class (ignoring numbers like "0" or "1" in the text)
    detected_name = class_names[index].lower()

    # 5. Corrected Logic for Color
    # This checks if the word "no" or "without" is in your label name
    if "no" in detected_name or "without" in detected_name:
        color = (0, 0, 255)  # RED
        status = "NO MASK"
    else:
        color = (0, 255, 0)  # GREEN
        status = "MASK DETECTED"

    # 6. UI Overlay
    # Only show the label if the AI is more than 70% confident
    if probability > 0.70:
        display_text = f"{status} ({int(probability*100)}%)"
        cv2.rectangle(display_img, (0,0), (w, h), color, 15)
        cv2.putText(display_img, display_text, (50, 80), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)
    else:
        cv2.putText(display_img, "Scanning...", (50, 80), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 2)

    cv2.imshow("M4 Face Mask Detector", display_img)

    if cv2.waitKey(1) == 27: break

camera.release()
cv2.destroyAllWindows()