feat: only copy images with confidence over 95%
Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>
This commit is contained in:
3
sort.py
3
sort.py
@@ -13,6 +13,7 @@ def sort_images():
|
||||
MODEL_PATH = 'garage_door_cnn.pth'
|
||||
SOURCE_DIR = 'data/hourly_photos/'
|
||||
DEST_DIR = 'data/sorted/open/'
|
||||
CONFIDENCE_THRESHOLD = 0.95 # Only copy if confidence is over this value
|
||||
|
||||
|
||||
# The classes are sorted alphabetically by ImageFolder: ['closed', 'open']
|
||||
@@ -63,7 +64,7 @@ def sort_images():
|
||||
probabilities = F.softmax(output, dim=1)
|
||||
confidence, pred_idx = torch.max(probabilities, 1)
|
||||
|
||||
if pred_idx.item() == TARGET_CLASS_IDX:
|
||||
if pred_idx.item() == TARGET_CLASS_IDX and confidence.item() > CONFIDENCE_THRESHOLD:
|
||||
print(f"Found 'open' image: {file_path} with confidence: {confidence.item():.4f}")
|
||||
# Copy file
|
||||
shutil.copy(file_path, os.path.join(DEST_DIR, filename))
|
||||
|
Reference in New Issue
Block a user