linefinder / Code:Scripts /batch_queue_direction.py
deansmile123's picture
Upload folder using huggingface_hub
b27cd24 verified
import os
import numpy as np
# Import the helpers from your existing script
from visual_3D import (
load_depth_npy,
load_orientations_npy,
bbox_centers_to_3d,
ransac_line_3d,
estimate_queue_forward_direction,
classify_queue_facing_direction,
)
from eval_gpt import load_all_json_recursive_with_paths # from your attached script
def load_image_paths(list_path: str):
paths = []
with open(list_path, "r", encoding="utf-8") as f:
for line in f:
p = line.strip()
if p and not p.startswith("#"):
paths.append(p)
return paths
IMAGE_LIST_PATH = "/scratch/ds5725/linefinder/LineFinder/Code:Scripts/olivia_luna_image_paths.txt"
# Find matching images
IMAGE_PATHS = load_image_paths(IMAGE_LIST_PATH)
print(f"Found {len(IMAGE_PATHS)} images that have GT JSONs.")
# for p in matched_images:
# print(p)
# Optional: save to file
# out_file = "images_with_gt.txt"
# with open(out_file, "w") as f:
# for p in matched_images:
# f.write(p + "\n")
# print(f"\nSaved list to {out_file}")
# exit()
# -------------------------------------------------------------------
# USER SETTINGS
# -------------------------------------------------------------------
# List of image paths to process (edit this list)
# IMAGE_PATHS = [
# r"D:\cv\LineFinder\queue_images\IMG_5954.JPG",
# r"D:\cv\LineFinder\queue_images\IMG_5955.JPG",
# r"D:\cv\LineFinder\queue_images\IMG_5956.JPG",
# ]
def load_focal_lengths(txt_path):
"""
Reads focal_length_px.txt and returns:
{ 'IMG_5954': 441.63, 'IMG_5955': 440.83, ... }
"""
focal_dict = {}
with open(txt_path, "r") as f:
for line in f:
if not line.strip():
continue
name, val = line.strip().split()
focal_dict[name] = float(val)
return focal_dict
# Where the depth .npy and bbox/orient .npy files live
DEPTH_DIR = "/scratch/ds5725/linefinder/LineFinder/depth_map"
BBOX_ORIENT_DIR = "/scratch/ds5725/linefinder/LineFinder/bbox_orient"
# Focal length in pixels (set to whatever you used in 3D_visual.py)
FOCAL_TXT = "/scratch/ds5725/linefinder/LineFinder/focal_length_px.txt"
focal_dict = load_focal_lengths(FOCAL_TXT)
print(f"Loaded {len(focal_dict)} focal-length records")
# RANSAC settings (same as your single-image script)
RANSAC_NUM_ITERS = 1000
RANSAC_DIST_THRESH = 0.8
RANSAC_MIN_INLIER_RATIO = 0.3
# front/back tolerance in degrees (same logic as before)
FRONT_BACK_TOL_DEG = 45.0
# -------------------------------------------------------------------
# HELPER: process one image
# -------------------------------------------------------------------
def classify_queue_for_image(
image_path: str,
depth_dir: str,
bbox_orient_dir: str,
f_px: float,
) -> str:
"""
Given an RGB image path, load the corresponding depth, bbox and
orientation npy files, run the 3D line + orientation logic, and
return the facing direction label (away / towards / sideways-* / unknown).
"""
if not os.path.isfile(image_path):
print(f"[WARN] Image not found: {image_path}")
return "missing-image"
# Get image ID from filename, e.g. IMG_5955 from IMG_5955.JPG
base = os.path.basename(image_path)
image_id, _ = os.path.splitext(base) # ("IMG_5955", ".JPG")
depth_npy_path = os.path.join(depth_dir, image_id + ".npy")
bboxes_npy_path = os.path.join(bbox_orient_dir, image_id + "_bboxes.npy")
orient_npy_path = os.path.join(bbox_orient_dir, image_id + "_orient.npy")
# Check all required files
missing_any = False
for p in [depth_npy_path, bboxes_npy_path, orient_npy_path]:
if not os.path.isfile(p):
print(f"[WARN] Missing file for {image_id}: {p}")
missing_any = True
if missing_any:
return "missing-data"
# ----- Load data -----
depth = load_depth_npy(depth_npy_path)
bboxes = np.load(bboxes_npy_path).astype(np.float32)
orientations_deg = load_orientations_npy(orient_npy_path)
# ----- 2D -> 3D -----
points_3d, centers_uv = bbox_centers_to_3d(bboxes, depth, f_px)
if points_3d.shape[0] < 2:
print(f"[INFO] Not enough 3D points to fit a line for {image_id}.")
return "too-few-points"
# ----- RANSAC line fit -----
line_point, line_dir, inlier_mask = ransac_line_3d(
points_3d,
num_iters=RANSAC_NUM_ITERS,
dist_thresh=RANSAC_DIST_THRESH,
min_inliers_ratio=RANSAC_MIN_INLIER_RATIO,
)
# ----- Use orientations to pick line direction -----
queue_forward_dir_3d, score = estimate_queue_forward_direction(
line_dir_3d=line_dir,
orientations_deg=orientations_deg,
inlier_mask=inlier_mask,
)
print(f"[DEBUG] {image_id}: avg alignment score = {score:.3f}")
# ----- Classify facing direction -----
facing_label = classify_queue_facing_direction(
queue_forward_dir_3d,
front_back_tolerance_deg=FRONT_BACK_TOL_DEG,
)
return facing_label
# -------------------------------------------------------------------
# MAIN: loop over all image paths
# -------------------------------------------------------------------
if __name__ == "__main__":
results = {}
for img_path in IMAGE_PATHS:
image_id = os.path.splitext(os.path.basename(img_path))[0]
# lookup focal length
if image_id not in focal_dict:
print(f"[WARN] No focal length found for {image_id}, skipping")
continue
label = classify_queue_for_image(
image_path=img_path,
depth_dir=DEPTH_DIR,
bbox_orient_dir=BBOX_ORIENT_DIR,
f_px=focal_dict[image_id],
)
results[img_path] = label
print(f"{img_path} -> {label}")
# (Optional) write to a CSV file
out_csv = "OL_queue_facing_results.csv"
try:
import csv
with open(out_csv, "w", newline="", encoding="utf-8") as f:
writer = csv.writer(f)
writer.writerow(["image_path", "facing_direction"])
for img_path, label in results.items():
writer.writerow([img_path, label])
print(f"\nSaved results to {out_csv}")
except Exception as e:
print(f"[WARN] Could not save CSV: {e}")