Spaces:
Running
Running
File size: 2,907 Bytes
5233510 |
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 |
import pandas as pd
import json
import re
# =========================
# EDIT THESE TWO VARIABLES
# =========================
INPUT_CSV = "raw_data.csv"
OUTPUT_CSV = "cleaned_data.csv"
def extract_video_id(filename: str):
"""'0111.mp4' -> '0111' (string, to match your existing style)"""
m = re.match(r"(\d+)", str(filename).strip())
return m.group(1) if m else None
def parse_video_labels(field):
"""
Parse Label Studio `videoLabels` JSON field.
Returns list of dicts: [{"label": str, "start": int, "end": int}, ...]
"""
if pd.isna(field):
return []
try:
items = json.loads(field)
except Exception:
return []
out = []
for entry in items if isinstance(items, list) else [items]:
labels = entry.get("timelinelabels") or entry.get("timelineLabels") or []
ranges = entry.get("ranges") or []
# normalize labels
if isinstance(labels, str):
labels = [labels]
labels = [str(x).strip() for x in labels if str(x).strip()]
for rr in ranges:
try:
start = int(rr.get("start"))
end = int(rr.get("end"))
except Exception:
continue
if start > end:
start, end = end, start
for lab in labels:
out.append({"label": lab, "start": start, "end": end})
return out
def main():
df = pd.read_csv(INPUT_CSV)
records = []
for _, row in df.iterrows():
filename = str(row.get("filename", "")).strip()
if not filename:
continue
video_id = extract_video_id(filename)
video_path = f"videos/{filename}"
labels = parse_video_labels(row.get("videoLabels", "[]"))
# If you want to DROP videos with no labels, replace this block with: `if not labels: continue`
if not labels:
records.append(
{
"filename": filename,
"video_id": video_id,
"video_path": video_path,
"label": None,
"start": None,
"end": None,
}
)
else:
for lab in labels:
records.append(
{
"filename": filename,
"video_id": video_id,
"video_path": video_path,
"label": lab["label"],
"start": lab["start"],
"end": lab["end"],
}
)
cleaned = pd.DataFrame(
records,
columns=["filename", "video_id", "video_path", "label", "start", "end"],
)
cleaned.to_csv(OUTPUT_CSV, index=False)
print(f"✅ Wrote {len(cleaned)} rows -> {OUTPUT_CSV}")
if __name__ == "__main__":
main()
|