updated requirements

requirements.txt

@@ -1,13 +1,5 @@
-numpy
-pandas
-matplotlib
-seaborn 
 tensorflow
-opencv-python
-fastapi
-uvicorn
-python-dotenv
 Pillow
-scikit-learn
-ultralytics 
-gradio
+numpy
+gradio
+huggingface_hub

src/pipeline/prediction_pipeline.py

@@ -5,12 +5,30 @@ from PIL import Image, ImageDraw, ImageFont
 import tensorflow as tf
 import json
 
+from huggingface_hub import hf_hub_download
+
 from src.config import Config
 
-CATEGORY_MODEL_PATH = Path(Config.MODEL_DIR) / "category_classifier.keras"
-FRESHNESS_MODEL_PATH = Path(Config.MODEL_DIR) / "mobilenetv2_baseline.keras"
+# ------------------------
+# Load models from Hugging Face Hub
+# ------------------------
+
+CATEGORY_MODEL_PATH = hf_hub_download(
+    repo_id="samithcs/Category_Classifier",
+    filename="category_classifier.keras"
+)
+
+FRESHNESS_MODEL_PATH = hf_hub_download(
+    repo_id="samithcs/Classifier_Fruits_Vegetables",
+    filename="mobilenetv2_baseline.keras"
+)
+
+# ------------------------
+# Load labels from LOCAL FILE
+# ------------------------
+
+CATEGORY_LABELS_PATH = Path("artifacts/models/category_labels.json")
 
-CATEGORY_LABELS_PATH = Path(Config.MODEL_DIR) / "category_labels.json"
 with open(CATEGORY_LABELS_PATH, "r") as f:
     CATEGORY_LABELS = json.load(f)
 
@@ -19,6 +37,7 @@ FRESHNESS_LABELS = ["Fresh", "Rotten"]
 CATEGORY_IMG_SIZE = (224, 224)
 FRESHNESS_IMG_SIZE = (224, 224)
 
+
 class PredictionPipeline:
     def __init__(self):
         self.category_model = tf.keras.models.load_model(CATEGORY_MODEL_PATH)
@@ -28,33 +47,29 @@ class PredictionPipeline:
 
     def _preprocess_image(self, img, target_size, normalize=True):
         if isinstance(img, (str, Path)):
-            img_path = Path(img)
-            if not img_path.exists():
-                raise FileNotFoundError(f"Image not found: {img_path.resolve()}")
-            img = Image.open(str(img_path)).convert("RGB")
+            img = Image.open(str(img)).convert("RGB")
         elif isinstance(img, np.ndarray):
-            if img.ndim == 3 and img.shape[2] == 3:
-                img = Image.fromarray(img)
-            else:
-                raise ValueError("NumPy input must be shape (H, W, 3)")
+            img = Image.fromarray(img)
+
         img_resized = img.resize(target_size)
         img_array = np.array(img_resized).astype("float32")
+
         if normalize:
             img_array = img_array / 255.0
+
         img_array = np.expand_dims(img_array, axis=0)
         return img_array, img
 
- 
     def predict(self, img):
-   
-        cat_img_array, pil_img = self._preprocess_image(img, CATEGORY_IMG_SIZE, normalize=True)
+        # Category prediction
+        cat_img_array, pil_img = self._preprocess_image(img, CATEGORY_IMG_SIZE)
         cat_pred = self.category_model.predict(cat_img_array)
         cat_idx = int(np.argmax(cat_pred))
         cat_label = self.category_labels[cat_idx]
         cat_score = float(np.max(cat_pred))
 
-      
-        fresh_img_array, _ = self._preprocess_image(img, FRESHNESS_IMG_SIZE, normalize=True)
+        # Freshness prediction
+        fresh_img_array, _ = self._preprocess_image(img, FRESHNESS_IMG_SIZE)
         fresh_pred = self.freshness_model.predict(fresh_img_array)
         fresh_idx = int(np.argmax(fresh_pred))
         fresh_label = self.freshness_labels[fresh_idx]
@@ -63,31 +78,24 @@ class PredictionPipeline:
         return {
             "category": {"label": cat_label, "idx": cat_idx, "score": cat_score},
             "freshness": {"label": fresh_label, "idx": fresh_idx, "score": fresh_score},
-            "pil_img": pil_img
+            "pil_img": pil_img,
         }
 
     def annotate(self, img, result, font_size=28):
         pil_img = result["pil_img"]
         draw = ImageDraw.Draw(pil_img)
-        text = f"{result['category']['label']} ({result['category']['score']:.2f}) | " \
-               f"{result['freshness']['label']} ({result['freshness']['score']:.2f})"
+
+        text = (
+            f"{result['category']['label']} ({result['category']['score']:.2f}) | "
+            f"{result['freshness']['label']} ({result['freshness']['score']:.2f})"
+        )
+
         try:
             font = ImageFont.truetype("arial.ttf", font_size)
         except:
             font = ImageFont.load_default()
-        draw.rectangle([0, 0, pil_img.width, font_size+8], fill=(0,0,0,160))
-        draw.text((5, 2), text, fill=(255, 255, 255), font=font)
-        return pil_img
 
-if __name__ == "__main__":
-    pipeline = PredictionPipeline()
-
-    img_path = r"artifacts\data\category\test\Apple\rottenApple (175).jpg"
-    
-    result = pipeline.predict(img_path)
-    
-    print("Prediction:", result["category"], "|", result["freshness"])
-
-    annotated_img = pipeline.annotate(img_path, result)
+        draw.rectangle([0, 0, pil_img.width, font_size + 8], fill=(0, 0, 0, 160))
+        draw.text((5, 2), text, fill=(255, 255, 255), font=font)
 
-    annotated_img.save(r"artifacts\results\annotated_result.jpg")
+        return pil_img