feat: save high-confidence images to sorted directories

Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>
This commit is contained in:
2025-08-01 16:30:02 +00:00
parent eb5f9ba00a
commit 20c433af2d

View File

@@ -9,7 +9,7 @@ import asyncio
import aiohttp import aiohttp
from aiohttp import web from aiohttp import web
import io import io
from datetime import datetime from datetime import datetime, timedelta
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
@@ -32,6 +32,7 @@ UNSURE_CONFIDENCE_THRESHOLD = 0.97
PREDICTION_HISTORY = [] PREDICTION_HISTORY = []
PREDICTION_HISTORY_MAX_LENGTH = 3 PREDICTION_HISTORY_MAX_LENGTH = 3
PREVIOUS_STATE = "unknown" PREVIOUS_STATE = "unknown"
LAST_OPEN_SAVE_TIME = None
# --- Model Inference --- # --- Model Inference ---
def get_prediction(model, image_bytes, device): def get_prediction(model, image_bytes, device):
@@ -61,6 +62,7 @@ def get_prediction(model, image_bytes, device):
# --- Background Task --- # --- Background Task ---
async def monitor_garage_door(app): async def monitor_garage_door(app):
"""Periodically fetches an image and logs the garage door status.""" """Periodically fetches an image and logs the garage door status."""
global LAST_OPEN_SAVE_TIME
logging.info("Starting garage door monitoring task.") logging.info("Starting garage door monitoring task.")
session = app['client_session'] session = app['client_session']
model = app['model'] model = app['model']
@@ -89,11 +91,10 @@ async def monitor_garage_door(app):
if len(PREDICTION_HISTORY) > PREDICTION_HISTORY_MAX_LENGTH: if len(PREDICTION_HISTORY) > PREDICTION_HISTORY_MAX_LENGTH:
PREDICTION_HISTORY.pop(0) PREDICTION_HISTORY.pop(0)
timestamp = datetime.now().isoformat().replace(':', '-')
filename = f"{timestamp}.jpg"
if confidence < UNSURE_CONFIDENCE_THRESHOLD: 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 # Construct path and save file
unsure_dir = os.path.join('data', 'unsure', prediction) unsure_dir = os.path.join('data', 'unsure', prediction)
os.makedirs(unsure_dir, exist_ok=True) os.makedirs(unsure_dir, exist_ok=True)
@@ -103,6 +104,29 @@ async def monitor_garage_door(app):
f.write(image_bytes) f.write(image_bytes)
logging.info(f"Low confidence prediction: {prediction} ({confidence:.4f}). Saved for review: {filepath}") 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: else:
logging.error(f"Failed to fetch image. Status: {response.status}, Reason: {response.reason}") logging.error(f"Failed to fetch image. Status: {response.status}, Reason: {response.reason}")