[FEAT]: add time and daily tags to select specific times on a daily basis

This commit is contained in:
2026-05-07 15:35:47 +02:00
parent 97d2896bfd
commit 025889728d

View File

@@ -128,11 +128,14 @@ def load_frames(zip_path: Path, max_frames: int):
# MAIN APP
# ─────────────────────────────────────────────
class Annotator(QMainWindow):
def __init__(self, data_dir: Path, out_dir: Path, clip: str = None):
def __init__(self, data_dir: Path, out_dir: Path, clip: str = None, target_time: str = None, daily: bool = False):
super().__init__()
self.data_dir = Path(data_dir)
self.out_dir = Path(out_dir)
self.target_time = target_time
self.daily = daily
self.current_date = None
self.history = []
self.erase_mode = False
@@ -161,9 +164,11 @@ class Annotator(QMainWindow):
df["datetime"] = df["filename"].apply(
lambda x: pd.to_datetime(x.stem.split("_")[1], errors="coerce")
)
# sort by datetime
df = df.sort_values("datetime").reset_index(drop=True)
return df
def _load_clip(self, specific: str = None):
def _load_clip(self, specific: str = None, next_day: bool = False):
if specific is not None:
matches = list(self.data_dir.glob(f"{specific}.zip"))
if not matches:
@@ -180,7 +185,63 @@ class Annotator(QMainWindow):
]
if not remaining:
raise RuntimeError("No remaining clips to annotate")
self.filename = np.random.choice(remaining)
if self.target_time or self.daily:
# Parse target time (format: HH:MM)
if self.target_time:
target_hour, target_minute = map(int, self.target_time.split(":"))
else:
target_hour, target_minute = 12, 0 # Default to noon
target_seconds = target_hour * 3600 + target_minute * 60
# Get datetimes for remaining files
remaining_datetimes = [
self.df[self.df["filename"] == f]["datetime"].values[0]
for f in remaining
]
# Group by day
df_remaining = pd.DataFrame({
"filename": remaining,
"datetime": remaining_datetimes
})
df_remaining["date"] = df_remaining["datetime"].dt.date
# In daily mode, filter to next day if needed
if self.daily and next_day and self.current_date is not None:
import datetime
next_date = self.current_date + datetime.timedelta(days=1)
df_remaining = df_remaining[df_remaining["date"] >= next_date]
if df_remaining.empty:
raise RuntimeError("No remaining clips to annotate")
# For each day, find the clip closest to target time
closest_clips = []
dates_list = []
for date, group in df_remaining.groupby("date"):
group = group.copy()
group["time_seconds"] = group["datetime"].dt.hour * 3600 + group["datetime"].dt.minute * 60
group["time_diff"] = (group["time_seconds"] - target_seconds).abs()
closest = group.loc[group["time_diff"].idxmin()]
closest_clips.append(closest["filename"])
dates_list.append(date)
# In daily mode, take only the first day's clip
if self.daily:
self.filename = closest_clips[0]
self.current_date = dates_list[0]
else:
# Take the first one (earliest by date/time)
self.filename = closest_clips[0]
self.current_date = dates_list[0]
else:
# take the earliest one (after sorting by datetime)
self.filename = remaining[0]
# Extract date from filename
import datetime
dt = self.df[self.df["filename"] == self.filename]["datetime"].values[0]
self.current_date = pd.Timestamp(dt).date()
self.frames, self.fps, self.dh, self.dw, self.h, self.w = load_frames(
self.filename, Config.MAX_FRAMES
@@ -489,7 +550,7 @@ class Annotator(QMainWindow):
def next_clip(self):
self.save()
self._load_clip()
self._load_clip(next_day=self.daily)
self.frame_i = 0
self.img.set_data(self.frames[0])
@@ -501,7 +562,7 @@ class Annotator(QMainWindow):
self._pending_answers = None
def skip_clip(self):
self._load_clip()
self._load_clip(next_day=self.daily)
self.frame_i = 0
self.img.set_data(self.frames[0])
@@ -518,9 +579,11 @@ class Annotator(QMainWindow):
# ─────────────────────────────────────────────
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--data", default="../torrent-flow/data/examples_for_annotations/")
parser.add_argument("--data", default="C:\\Users\\BONVALOT\\Documents\\Hydroscan\\data\\filtered_s3\\all_filtered_data")
parser.add_argument("--out", default="data/annotation_results/")
parser.add_argument("--clip", default=None, help="Stem name of a specific clip to load (e.g. 'left_20230501')")
parser.add_argument("--time", default=None, help="Target time to filter clips by day (format: HH:MM, e.g. '14:30'). Selects the closest clip to this time for each day.")
parser.add_argument("--daily", action="store_true", help="Load only 1 clip per day at the specified time (requires --time).")
return parser.parse_args()
@@ -529,7 +592,7 @@ if __name__ == "__main__":
app = QApplication([])
win = Annotator(Path(args.data), Path(args.out), clip=args.clip)
win = Annotator(Path(args.data), Path(args.out), clip=args.clip, target_time=args.time, daily=args.daily)
win.show()
app.exec()