feat: process images in random order and show progress
Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>
This commit is contained in:
13
sort.py
13
sort.py
@@ -4,6 +4,7 @@ from torchvision import transforms
|
|||||||
from PIL import Image
|
from PIL import Image
|
||||||
import os
|
import os
|
||||||
import shutil
|
import shutil
|
||||||
|
import random
|
||||||
|
|
||||||
from model import (CropLowerRightTriangle, GarageDoorCNN, TRIANGLE_CROP_WIDTH,
|
from model import (CropLowerRightTriangle, GarageDoorCNN, TRIANGLE_CROP_WIDTH,
|
||||||
TRIANGLE_CROP_HEIGHT, RESIZE_DIM)
|
TRIANGLE_CROP_HEIGHT, RESIZE_DIM)
|
||||||
@@ -45,10 +46,15 @@ def sort_images():
|
|||||||
|
|
||||||
# --- Process Images ---
|
# --- Process Images ---
|
||||||
print(f"Scanning images in {SOURCE_DIR}...")
|
print(f"Scanning images in {SOURCE_DIR}...")
|
||||||
|
|
||||||
|
# Get and shuffle filenames
|
||||||
|
filenames = [f for f in os.listdir(SOURCE_DIR) if os.path.isfile(os.path.join(SOURCE_DIR, f))]
|
||||||
|
random.shuffle(filenames)
|
||||||
|
total_files = len(filenames)
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
for filename in os.listdir(SOURCE_DIR):
|
for i, filename in enumerate(filenames):
|
||||||
file_path = os.path.join(SOURCE_DIR, filename)
|
file_path = os.path.join(SOURCE_DIR, filename)
|
||||||
if os.path.isfile(file_path):
|
|
||||||
try:
|
try:
|
||||||
image = Image.open(file_path).convert('RGB')
|
image = Image.open(file_path).convert('RGB')
|
||||||
|
|
||||||
@@ -65,7 +71,8 @@ def sort_images():
|
|||||||
confidence, pred_idx = torch.max(probabilities, 1)
|
confidence, pred_idx = torch.max(probabilities, 1)
|
||||||
|
|
||||||
if pred_idx.item() == TARGET_CLASS_IDX and confidence.item() > CONFIDENCE_THRESHOLD:
|
if pred_idx.item() == TARGET_CLASS_IDX and confidence.item() > CONFIDENCE_THRESHOLD:
|
||||||
print(f"Found 'open' image: {file_path} with confidence: {confidence.item():.4f}")
|
progress = f"({i + 1}/{total_files} - {(i + 1) / total_files * 100:.1f}%)"
|
||||||
|
print(f"{progress} Found 'open' image: {file_path} with confidence: {confidence.item():.4f}")
|
||||||
# Copy file
|
# Copy file
|
||||||
shutil.copy(file_path, os.path.join(DEST_DIR, filename))
|
shutil.copy(file_path, os.path.join(DEST_DIR, filename))
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user