Spaces:
Sleeping
Sleeping
| 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" |