diff --git a/src/river_annotation_tool/annotation_script.py b/src/river_annotation_tool/annotation_script.py index cca63b8..3f9c6b1 100644 --- a/src/river_annotation_tool/annotation_script.py +++ b/src/river_annotation_tool/annotation_script.py @@ -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()