Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import traceback | |
| import numpy as np | |
| from PIL import Image | |
| from tensorflow.keras.models import load_model | |
| from tensorflow.keras.applications.efficientnet import preprocess_input | |
| BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| MODEL_PATH = os.path.join(BASE_DIR, 'mission_model.h5') | |
| LABELS_PATH = os.path.join(BASE_DIR, 'labels.txt') | |
| class Predictor: | |
| def __init__(self): | |
| self.model = None | |
| self.class_names = [] | |
| self._load_model() | |
| def _load_model(self): | |
| print("🧠 Loading TensorFlow CNN Brain...") | |
| if not os.path.exists(MODEL_PATH): | |
| print(f"❌ ERROR: {MODEL_PATH} not found. You need to train the model!") | |
| return | |
| try: | |
| self.model = load_model(MODEL_PATH) | |
| print("✅ Model loaded successfully!") | |
| except Exception as e: | |
| print(f"❌ Failed to load model: {e}") | |
| # Load Labels | |
| try: | |
| with open(LABELS_PATH, 'r') as f: | |
| self.class_names = [line.strip() for line in f.readlines()] | |
| print(f"🏷️ Labels loaded: {self.class_names}") | |
| except FileNotFoundError: | |
| print("❌ ERROR: labels.txt not found.") | |
| self.class_names = [] | |
| # 🔥 WARMUP STEP (Optimization) | |
| if self.model: | |
| print("🔥 Warming up model for instant first-prediction...") | |
| dummy_image = np.zeros((1, 224, 224, 3), dtype=np.float32) | |
| self.model.predict(dummy_image, verbose=0) | |
| print("⚡ AI is fully optimized and ready!") | |
| def predict(self, file_bytes): | |
| """ | |
| Runs the image through the custom EfficientNet CNN. | |
| """ | |
| if not self.model or not self.class_names: | |
| return {"category": "Non_SDG_Invalid", "confidence": 0, "reason": "Model offline or missing."} | |
| try: | |
| # 1. Read image using PIL (just like in train_ai.py) | |
| img = Image.open(io.BytesIO(file_bytes)).convert('RGB') | |
| # 2. Resize to 224x224 (EfficientNetB0 input size) | |
| img = img.resize((224, 224), Image.LANCZOS) | |
| # 3. Apply EfficientNetB0 preprocess_input | |
| img_array = np.array(img, dtype=np.float32) | |
| img_array = preprocess_input(img_array) | |
| img_array = np.expand_dims(img_array, axis=0) | |
| # 4. Predict | |
| predictions = self.model.predict(img_array) | |
| score = predictions[0] | |
| top_index = np.argmax(score) | |
| label = self.class_names[top_index] | |
| confidence = int(np.max(score) * 100) | |
| # Clean up label if it has the SDG prefix (e.g. SDG12_Recycling -> Recycling) | |
| # The verdict.py MISSION_MAP expects "Recycling", "Planting", etc. | |
| category = label | |
| if "_" in label and label.startswith("SDG"): | |
| # E.g. "SDG12_Recycling" -> "Recycling" | |
| category = label.split("_", 1)[1] | |
| # If there are multiple underscores (like SDG13_15_Planting), take the last part | |
| if "_" in category: | |
| category = category.rsplit("_", 1)[-1] | |
| elif label == "Non_SDG_Invalid": | |
| category = "Non_SDG_Invalid" | |
| # Quick check for combined strings | |
| if "Planting" in label: category = "Planting" | |
| if "Cleanup" in label: category = "Cleanup" | |
| if "Donation" in label: category = "Donation" | |
| if "Cities" in label or "Sustainable" in label: category = "Sustainable_Cities" | |
| if "Local" in label: category = "Support_Local" | |
| if "Health" in label: category = "Health" | |
| if "Energy" in label: category = "Energy" | |
| if "Education" in label: category = "Education" | |
| return { | |
| "category": category, | |
| "confidence": confidence, | |
| "reason": f"Predicted {label} with {confidence}% confidence" | |
| } | |
| except Exception as e: | |
| traceback.print_exc() | |
| print(f"⚠️ Predictor error: {e}") | |
| return {"category": "Non_SDG_Invalid", "confidence": 0, "reason": str(e)} | |
| def get_model_name(self): | |
| return "Custom CNN (mission_model.h5)" | |