No public description

PiperOrigin-RevId: 814387436

Author: tensorflower-gardener

official/projects/waste_identification_ml/llm_applications/milk_pouch_detection/classify_images.py

@@ -0,0 +1,60 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Classifies images based on Image Classifier."""
+
+import glob
+import os
+import shutil
+
+from absl import flags
+import torch
+
+from official.projects.waste_identification_ml.llm_applications.milk_pouch_detection import models
+
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+ "input_dir", "tempdir", "Directory containing image frames to process."
+)
+
+
+# Path to the custom trained model for Image Classifier.
+IMAGE_CLASSIFIER_WEIGHTS = (
+ "milk_pouch_project/image_classifier_model/best_vit_model_epoch_131.pt"
+)
+CLASS_NAMES = ["dairy", "other"]
+
+
+def main(_) -> None:
+ os.makedirs("dairy_packets", exist_ok=True)
+ os.makedirs("others", exist_ok=True)
+
+ classifier = models.ImageClassifier(
+ model_path=IMAGE_CLASSIFIER_WEIGHTS,
+ class_names=CLASS_NAMES,
+ device="cuda" if torch.cuda.is_available() else "cpu",
+ )
+
+ files = glob.glob(os.path.join(FLAGS.input_dir, "*"))
+ print(f"Found {len(files)} images to process...")
+
+ total_dairy_packets = 0
+ for path in files:
+ pred_class, _ = classifier.classify(path)
+ if pred_class == "dairy":
+ total_dairy_packets += 1
+ shutil.move(path, os.path.join("dairy_packets", os.path.basename(path)))
+ else:
+ shutil.move(path, os.path.join("others", os.path.basename(path)))

official/projects/waste_identification_ml/llm_applications/milk_pouch_detection/run_pipeline.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# --- Parse command-line arguments ---
+while [[ "$#" -gt 0 ]]; do
+ case $1 in
+ --input_dir=*) input_dir="${1#*=}"; shift ;;
+ --category_name=*) category_name="${1#*=}"; shift ;;
+ *) echo "❌ Unknown parameter passed: $1"; exit 1 ;;
+ esac
+done
+
+# --- Check if required arguments were provided ---
+if [ -z "$input_dir" ]; then
+ echo "❌ Error: --input_dir must be specified"
+ echo "✅ Usage: ./run_pipeline.sh --input_dir=/path/to/images --category_name=category"
+ exit 1
+fi
+
+if [ -z "$category_name" ]; then
+ echo "❌ Error: --category_name must be specified"
+ echo "✅ Usage: ./run_pipeline.sh --input_dir=/path/to/images --category_name=category"
+ exit 1
+fi
+
+# --- Run pipeline ---
+echo "✅ Activating virtual environment..."
+source myenv/bin/activate
+
+echo "🚀 Running detect_and_segment.py..."
+echo " Input directory: $input_dir"
+echo " Category name: $category_name"
+python3 detect_and_segment.py --input_dir="$input_dir" --category_name="$category_name"
+
+echo "🧠 Running classify.py..."
+python3 classify.py
+
+echo "🧹 Deactivating virtual environment..."
+deactivate
+
+echo "✅ Done."