diff --git a/server.py b/server.py index 4104f5c..c195ee8 100644 --- a/server.py +++ b/server.py @@ -4,6 +4,7 @@ from aiohttp import web import logging import os import io +from datetime import datetime import torch import torch.nn.functional as F @@ -25,6 +26,7 @@ MODEL_PATH = 'garage_door_cnn.pth' CLASS_NAMES = ['closed', 'open'] # From training, sorted alphabetically POLL_INTERVAL_SECONDS = 10 REQUEST_TIMEOUT_SECONDS = 5 +UNSURE_CONFIDENCE_THRESHOLD = 0.97 # --- Model Inference --- def get_prediction(model, image_bytes, device): @@ -71,6 +73,21 @@ async def monitor_garage_door(app): if result: prediction, confidence = result logging.debug(f"Garage door status: {prediction} (confidence: {confidence:.4f})") + + if confidence < UNSURE_CONFIDENCE_THRESHOLD: + # Sanitize timestamp for use in filename + timestamp = datetime.now().isoformat().replace(':', '-') + filename = f"{timestamp}.jpg" + + # Construct path and save file + unsure_dir = os.path.join('data', 'unsure', prediction) + os.makedirs(unsure_dir, exist_ok=True) + filepath = os.path.join(unsure_dir, filename) + + with open(filepath, 'wb') as f: + f.write(image_bytes) + + logging.info(f"Low confidence prediction: {prediction} ({confidence:.4f}). Saved for review: {filepath}") else: logging.error(f"Failed to fetch image. Status: {response.status}, Reason: {response.reason}")