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:
34
server.py
34
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}")
|
||||
|
||||
|
Reference in New Issue
Block a user