""" Wildfire Detection Inference Script ==================================== For LOCAL TESTING: Uses trained RL model (no API key required) For SUBMISSION: Uses OpenAI client with API key injected by judges """ import os import sys import base64 import pickle from io import BytesIO from typing import List, Optional, Dict, Any from collections import defaultdict from openai import OpenAI import numpy as np from PIL import Image # Environment variables API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-3B-Instruct") HF_TOKEN = os.getenv("HF_TOKEN") # Optional - if you use from_docker_image(): LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") API_KEY = os.getenv("API_KEY", HF_TOKEN) # Check if we can use LLM mode USE_LLM = API_KEY is not None if USE_LLM: client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) else: client = None print("NOTE: Running in LOCAL MODE (no API_KEY). Using trained RL model.") # RL Model class for local mode class TrainedQLearningAgent: """Trained Q-learning agent for local inference.""" def __init__(self): model_path = os.path.join(os.path.dirname(__file__), "q_model.pkl") if os.path.exists(model_path): with open(model_path, "rb") as f: self.q_table = defaultdict(lambda: [0.0, 0.0, 0.0, 0.0], pickle.load(f)) print(f"Loaded RL model from: {model_path}") else: print("WARNING: No trained model found, using fallback!") self.q_table = defaultdict(lambda: [0.0, 0.0, 0.0, 0.0]) self.actions = ["Alert", "Scan", "Ignore", "Deploy"] def get_state_key(self, obs): prediction = obs.get("prediction", np.zeros(3)) if hasattr(prediction, "tolist"): p = prediction.tolist() else: p = list(prediction) fire_high = 0 if p[0] < 0.5 else 1 smoke_high = 0 if p[1] < 0.5 else 1 return (fire_high, smoke_high) def choose_action(self, state): q_values = self.q_table.get(state, [0.0, 0.0, 0.0, 0.0]) if hasattr(q_values, "tolist"): q_values = q_values.tolist() return q_values.index(max(q_values)) def get_action_name(self, idx): return self.actions[idx] # Import environment sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from environments.wildfire_detection.wildfire_env import WildfireDetectionEnv def encode_image_to_base64(image_array: np.ndarray) -> str: """Convert numpy image to base64 data URI.""" image = Image.fromarray(image_array) buffer = BytesIO() image.save(buffer, format="JPEG", quality=85) buffer.seek(0) data = base64.b64encode(buffer.read()).decode("utf-8") return f"data:image/jpeg;base64,{data}" def build_observation_prompt(step: int, obs: Dict[str, Any]) -> str: """Build user prompt from observation.""" prediction = obs.get("prediction", {}) fire_prob = prediction[0] if hasattr(prediction, "__getitem__") else 0 smoke_prob = prediction[1] if hasattr(prediction, "__getitem__") else 0 no_fire_prob = prediction[2] if hasattr(prediction, "__getitem__") else 0 gradcam = obs.get("gradcam_summary", "N/A") frame_id = ( obs.get("frame_id", [0])[0] if hasattr(obs.get("frame_id", [0]), "__len__") else 0 ) prompt = f"""Frame {frame_id} - Step {step} Prediction probabilities: - Fire: {fire_prob:.2%} - Smoke: {smoke_prob:.2%} - No Fire: {no_fire_prob:.2%} Grad-CAM analysis: {gradcam} What is your action? (Alert/Scan/Ignore/Deploy) """ return prompt.strip() def parse_action(response_text: str) -> str: """Parse LLM response to extract action.""" if not response_text: return "Scan" response = response_text.strip().lower() if "alert" in response: return "Alert" elif "deploy" in response: return "Deploy" elif "ignore" in response: return "Ignore" return "Scan" def action_to_index(action: str) -> int: actions = ["Alert", "Scan", "Ignore", "Deploy"] return actions.index(action) if action in actions else 1 MAX_STEPS = 20 TEMPERATURE = 0.2 MAX_TOKENS = 500 SYSTEM_PROMPT = """You are an autonomous wildfire detection agent. You control a drone monitoring for fires and smoke. Given camera frame observations with prediction probabilities and Grad-CAM analysis, decide the best action. Actions available: - Alert: Raise immediate alert if fire/smoke detected with high confidence - Scan: Continue scanning/observing the current frame more carefully - Ignore: No action needed, clear conditions - Deploy: Deploy resources to the detected location Output ONLY the action name: Alert, Scan, Ignore, or Deploy Consider: - Prediction confidence levels - Grad-CAM heatmap alignment (strong hotspot = more reliable detection) - Balance early detection vs false alarms - Sequential frame context """ def run_episode(env) -> Dict[str, Any]: """Run a single episode and return results.""" rewards_list = [] steps_executed = 0 # Use appropriate model name based on mode model_used = MODEL_NAME if USE_LLM else "trained_q_learning" print(f"[START] task=wildfire_detection env=wildfire_detection model={model_used}") # Initialize RL agent for local mode rl_agent = TrainedQLearningAgent() if not USE_LLM else None obs = env.reset() for step in range(1, MAX_STEPS + 1): steps_executed = step if USE_LLM: # LLM mode user_prompt = build_observation_prompt(step, obs) image_b64 = "" if "image" in obs and obs["image"] is not None: image_b64 = encode_image_to_base64(obs["image"]) messages = [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": [ {"type": "text", "text": user_prompt}, {"type": "image_url", "image_url": {"url": image_b64}} if image_b64 else {"type": "text", "text": "(no image)"}, ], }, ] try: completion = client.chat.completions.create( model=MODEL_NAME, messages=messages, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False, ) response_text = completion.choices[0].message.content or "" except Exception as e: print(f"Model request failed: {e}. Using fallback.") response_text = "Scan" action_str = parse_action(response_text) else: # RL mode state = rl_agent.get_state_key(obs) action_idx = rl_agent.choose_action(state) action_str = rl_agent.get_action_name(action_idx) action_idx = action_to_index(action_str) obs, reward, done, info = env.step(action_idx) rewards_list.append(reward) error_msg = "null" print( f"[STEP] step={step} action={action_str} reward={reward:.2f} done={str(done).lower()} error={error_msg}" ) if done: break total_reward = sum(rewards_list) rewards_str = ",".join(f"{r:.2f}" for r in rewards_list) success = total_reward > 0 print( f"[END] success={str(success).lower()} steps={steps_executed} rewards={rewards_str}" ) return { "success": success, "steps": steps_executed, "rewards": rewards_list, "total_reward": total_reward, } def main(): """Main entry point.""" model_path = os.path.join( os.path.dirname(__file__), "FirenetCNN1.h5", ) env = WildfireDetectionEnv(model_path=model_path) try: result = run_episode(env) print( f"Episode completed: success={result['success']}, steps={result['steps']}, total_reward={result['total_reward']:.2f}" ) finally: env.close() if __name__ == "__main__": main()