import cv2
import os
import mediapipe as mp

# --- Initialize MediaPipe Image Classifier ---
BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# ✅ Load model
model_path = "classifier.tflite"
# model_path = "2.tflite"
# https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path=model_path),
    max_results=3
)
classifier = ImageClassifier.create_from_options(options)

# --- Load images from folder ---
folder = r"C:\Users\R c\PycharmProjects\BG_Remover\images"
images = [os.path.join(folder, f) for f in os.listdir(folder)
          if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

if not images:
    print("❌ No images found in folder.")
    exit()

index = 0

# --- Main Loop ---
while True:
    image_path = images[index]
    frame = cv2.imread(image_path)

    if frame is None:
        print(f"⚠️ Skipping unreadable image: {image_path}")
        index = (index + 1) % len(images)
        continue

    # ✅ Scale down image (20% of original size)
    frame = cv2.resize(frame, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)

    # Convert to RGB for MediaPipe
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)

    # Classify the image
    result = classifier.classify(mp_image)

    # --- Get top label and score ---
    if result.classifications:
        category = result.classifications[0].categories[0]
        label = category.category_name
        score = category.score
        text = f"{label} ({score:.2f})"
    else:
        text = "No classification"

    # Draw label on image
    cv2.putText(frame, text, (20, 40),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

    cv2.imshow("E-Commerce Image Classification", frame)

    key = cv2.waitKey(0) & 0xFF
    if key == 27:  # ESC → exit
        break
    elif key == 32:  # SPACE → next image
        index = (index + 1) % len(images)

cv2.destroyAllWindows()


# put this as deafult address r"C:\Users\R c\PycharmProjects\BG_Remover\images"