[FEAT]: add time and daily tags to select specific times on a daily basis
This commit is contained in:
@@ -128,11 +128,14 @@ def load_frames(zip_path: Path, max_frames: int):
|
|||||||
# MAIN APP
|
# MAIN APP
|
||||||
# ─────────────────────────────────────────────
|
# ─────────────────────────────────────────────
|
||||||
class Annotator(QMainWindow):
|
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__()
|
super().__init__()
|
||||||
|
|
||||||
self.data_dir = Path(data_dir)
|
self.data_dir = Path(data_dir)
|
||||||
self.out_dir = Path(out_dir)
|
self.out_dir = Path(out_dir)
|
||||||
|
self.target_time = target_time
|
||||||
|
self.daily = daily
|
||||||
|
self.current_date = None
|
||||||
|
|
||||||
self.history = []
|
self.history = []
|
||||||
self.erase_mode = False
|
self.erase_mode = False
|
||||||
@@ -161,9 +164,11 @@ class Annotator(QMainWindow):
|
|||||||
df["datetime"] = df["filename"].apply(
|
df["datetime"] = df["filename"].apply(
|
||||||
lambda x: pd.to_datetime(x.stem.split("_")[1], errors="coerce")
|
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
|
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:
|
if specific is not None:
|
||||||
matches = list(self.data_dir.glob(f"{specific}.zip"))
|
matches = list(self.data_dir.glob(f"{specific}.zip"))
|
||||||
if not matches:
|
if not matches:
|
||||||
@@ -180,7 +185,63 @@ class Annotator(QMainWindow):
|
|||||||
]
|
]
|
||||||
if not remaining:
|
if not remaining:
|
||||||
raise RuntimeError("No remaining clips to annotate")
|
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.frames, self.fps, self.dh, self.dw, self.h, self.w = load_frames(
|
||||||
self.filename, Config.MAX_FRAMES
|
self.filename, Config.MAX_FRAMES
|
||||||
@@ -489,7 +550,7 @@ class Annotator(QMainWindow):
|
|||||||
|
|
||||||
def next_clip(self):
|
def next_clip(self):
|
||||||
self.save()
|
self.save()
|
||||||
self._load_clip()
|
self._load_clip(next_day=self.daily)
|
||||||
|
|
||||||
self.frame_i = 0
|
self.frame_i = 0
|
||||||
self.img.set_data(self.frames[0])
|
self.img.set_data(self.frames[0])
|
||||||
@@ -501,7 +562,7 @@ class Annotator(QMainWindow):
|
|||||||
self._pending_answers = None
|
self._pending_answers = None
|
||||||
|
|
||||||
def skip_clip(self):
|
def skip_clip(self):
|
||||||
self._load_clip()
|
self._load_clip(next_day=self.daily)
|
||||||
|
|
||||||
self.frame_i = 0
|
self.frame_i = 0
|
||||||
self.img.set_data(self.frames[0])
|
self.img.set_data(self.frames[0])
|
||||||
@@ -518,9 +579,11 @@ class Annotator(QMainWindow):
|
|||||||
# ─────────────────────────────────────────────
|
# ─────────────────────────────────────────────
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = argparse.ArgumentParser()
|
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("--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("--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()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
@@ -529,7 +592,7 @@ if __name__ == "__main__":
|
|||||||
|
|
||||||
app = QApplication([])
|
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()
|
win.show()
|
||||||
|
|
||||||
app.exec()
|
app.exec()
|
||||||
|
|||||||
Reference in New Issue
Block a user