diff --git a/server.py b/server.py index f221370..7976759 100644 --- a/server.py +++ b/server.py @@ -9,7 +9,7 @@ import asyncio import aiohttp from aiohttp import web import io -from datetime import datetime +from datetime import datetime, timedelta import torch import torch.nn.functional as F @@ -32,6 +32,7 @@ UNSURE_CONFIDENCE_THRESHOLD = 0.97 PREDICTION_HISTORY = [] PREDICTION_HISTORY_MAX_LENGTH = 3 PREVIOUS_STATE = "unknown" +LAST_OPEN_SAVE_TIME = None # --- Model Inference --- def get_prediction(model, image_bytes, device): @@ -61,6 +62,7 @@ def get_prediction(model, image_bytes, device): # --- Background Task --- async def monitor_garage_door(app): """Periodically fetches an image and logs the garage door status.""" + global LAST_OPEN_SAVE_TIME logging.info("Starting garage door monitoring task.") session = app['client_session'] model = app['model'] @@ -89,11 +91,10 @@ async def monitor_garage_door(app): if len(PREDICTION_HISTORY) > PREDICTION_HISTORY_MAX_LENGTH: PREDICTION_HISTORY.pop(0) + timestamp = datetime.now().isoformat().replace(':', '-') + filename = f"{timestamp}.jpg" + 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) @@ -103,6 +104,29 @@ async def monitor_garage_door(app): f.write(image_bytes) logging.info(f"Low confidence prediction: {prediction} ({confidence:.4f}). Saved for review: {filepath}") + else: + # High confidence, save to sorted + if prediction == 'open': + if LAST_OPEN_SAVE_TIME is None or (datetime.now() - LAST_OPEN_SAVE_TIME) > timedelta(minutes=5): + sorted_dir = os.path.join('data', 'sorted', 'open') + os.makedirs(sorted_dir, exist_ok=True) + filepath = os.path.join(sorted_dir, filename) + with open(filepath, 'wb') as f: + f.write(image_bytes) + LAST_OPEN_SAVE_TIME = datetime.now() + logging.info(f"Saved high-confidence 'open' image: {filepath}") + elif prediction == 'closed': + open_dir = os.path.join('data', 'sorted', 'open') + closed_dir = os.path.join('data', 'sorted', 'closed') + os.makedirs(open_dir, exist_ok=True) + os.makedirs(closed_dir, exist_ok=True) + num_open = len(os.listdir(open_dir)) + num_closed = len(os.listdir(closed_dir)) + if num_closed < num_open: + filepath = os.path.join(closed_dir, filename) + with open(filepath, 'wb') as f: + f.write(image_bytes) + logging.info(f"Saved high-confidence 'closed' image: {filepath}") else: logging.error(f"Failed to fetch image. Status: {response.status}, Reason: {response.reason}")