File size: 6,314 Bytes
b27cd24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
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}")