image-classifier / main.py
RaheelShah's picture
deployment_push
6538f2b
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"