Factor Studios commited on
Commit
eb80f12
·
verified ·
1 Parent(s): b6fb585

Update vision_analyzer.py

Browse files
Files changed (1) hide show
  1. vision_analyzer.py +439 -692
vision_analyzer.py CHANGED
@@ -1,692 +1,439 @@
1
- import os
2
- import json
3
- import requests
4
- import subprocess
5
- import shutil
6
- import time
7
- import re
8
- import threading
9
- from typing import Dict, List, Set, Optional
10
- from huggingface_hub import HfApi, list_repo_files
11
- from fastapi import FastAPI, File, UploadFile, Form
12
- from fastapi.responses import JSONResponse
13
- from pathlib import Path
14
- import smtplib
15
- from email.message import EmailMessage
16
- import tempfile
17
- import rarfile
18
- import zipfile
19
- import cv2
20
- import numpy as np
21
- from PIL import Image
22
- import torch
23
- from transformers import AutoProcessor, AutoModelForCausalLM
24
-
25
- # Initialize FastAPI
26
- app = FastAPI()
27
-
28
- # ==== CONFIGURATION ====
29
- HF_TOKEN = os.getenv("HF_TOKEN", "")
30
- SOURCE_REPO_ID = os.getenv("SOURCE_REPO", "Fred808/BG1")
31
-
32
- # Path Configuration
33
- DOWNLOAD_FOLDER = "downloads"
34
- EXTRACT_FOLDER = "extracted"
35
- FRAMES_OUTPUT_FOLDER = "extracted_frames"
36
- ANALYSIS_OUTPUT_FOLDER = "analysis_results"
37
-
38
- os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)
39
- os.makedirs(EXTRACT_FOLDER, exist_ok=True)
40
- os.makedirs(FRAMES_OUTPUT_FOLDER, exist_ok=True)
41
- os.makedirs(ANALYSIS_OUTPUT_FOLDER, exist_ok=True)
42
-
43
- # State Files
44
- DOWNLOAD_STATE_FILE = "download_progress.json"
45
- PROCESS_STATE_FILE = "process_progress.json"
46
- FAILED_FILES_LOG = "failed_files.log"
47
-
48
- # Processing Parameters
49
- CHUNK_SIZE = 1
50
- PROCESSING_DELAY = 2
51
- MAX_RETRIES = 3
52
- MIN_FREE_SPACE_GB = 2 # Minimum free space in GB before processing
53
-
54
- # Frame Extraction Parameters
55
- DEFAULT_FPS = 3 # Default frames per second for extraction
56
-
57
- # Initialize HF API
58
- hf_api = HfApi(token=HF_TOKEN)
59
-
60
- # Global State
61
- processing_status = {
62
- "is_running": False,
63
- "current_file": None,
64
- "total_files": 0,
65
- "processed_files": 0,
66
- "failed_files": 0,
67
- "extracted_courses": 0,
68
- "extracted_videos": 0,
69
- "extracted_frames_count": 0,
70
- "analyzed_frames_count": 0,
71
- "last_update": None,
72
- "logs": []
73
- }
74
-
75
- import torch
76
- import subprocess
77
- import sys
78
-
79
-
80
-
81
- device = "cpu" # Explicitly ensure CPU usage
82
-
83
- try:
84
- # Load the model, forcing the 'eager' (CPU-compatible) attention implementation
85
- vision_language_model_large = AutoModelForCausalLM.from_pretrained(
86
- "microsoft/Florence-2-Base",
87
- trust_remote_code=True
88
- ).to(device).eval()
89
- vision_language_processor_large = AutoProcessor.from_pretrained(
90
- "microsoft/Florence-2-Base",
91
- trust_remote_code=True
92
- )
93
- print("Florence-2 large model and processor loaded successfully on CPU using eager attention.")
94
- except Exception as e:
95
- print(f"Error loading Florence-2 model on CPU: {e}")
96
- print("Please ensure you have enough RAM and a compatible PyTorch version.")
97
- vision_language_model_large = None
98
- vision_language_processor_large = None
99
-
100
-
101
- def log_message(message: str):
102
- """Log messages with timestamp"""
103
- timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
104
- log_entry = f"[{timestamp}] {message}"
105
- print(log_entry)
106
- processing_status["logs"].append(log_entry)
107
- processing_status["last_update"] = timestamp
108
- if len(processing_status["logs"]) > 100:
109
- processing_status["logs"] = processing_status["logs"][-100:]
110
-
111
- def log_failed_file(filename: str, error: str):
112
- """Log failed files to persistent file"""
113
- with open(FAILED_FILES_LOG, "a") as f:
114
- f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {filename}: {error}\n")
115
-
116
- def get_disk_usage(path: str) -> Dict[str, float]:
117
- """Get disk usage statistics in GB"""
118
- statvfs = os.statvfs(path)
119
- total = statvfs.f_frsize * statvfs.f_blocks / (1024**3)
120
- free = statvfs.f_frsize * statvfs.f_bavail / (1024**3)
121
- used = total - free
122
- return {"total": total, "free": free, "used": used}
123
-
124
- def check_disk_space(path: str = ".") -> bool:
125
- """Check if there's enough disk space"""
126
- disk_info = get_disk_usage(path)
127
- if disk_info["free"] < MIN_FREE_SPACE_GB:
128
- log_message(f'⚠️ Low disk space: {disk_info["free"]:.2f}GB free, {disk_info["used"]:.2f}GB used')
129
- return False
130
- return True
131
-
132
- def cleanup_temp_files():
133
- """Clean up temporary files to free space"""
134
- log_message("🧹 Cleaning up temporary files...")
135
-
136
- # Clean old downloads (keep only current processing file)
137
- current_file = processing_status.get("current_file")
138
- for file in os.listdir(DOWNLOAD_FOLDER):
139
- if file != current_file and file.endswith((".rar", ".zip")):
140
- try:
141
- os.remove(os.path.join(DOWNLOAD_FOLDER, file))
142
- log_message(f"🗑️ Removed old download: {file}")
143
- except:
144
- pass
145
-
146
- def load_json_state(file_path: str, default_value):
147
- """Load state from JSON file"""
148
- if os.path.exists(file_path):
149
- try:
150
- with open(file_path, "r") as f:
151
- return json.load(f)
152
- except json.JSONDecodeError:
153
- log_message(f"⚠️ Corrupted state file: {file_path}")
154
- return default_value
155
-
156
- def save_json_state(file_path: str, data):
157
- """Save state to JSON file"""
158
- with open(file_path, "w") as f:
159
- json.dump(data, f, indent=2)
160
-
161
- def download_with_retry(url: str, dest_path: str, max_retries: int = 3) -> bool:
162
- """Download file with retry logic and disk space checking"""
163
- if not check_disk_space():
164
- cleanup_temp_files()
165
- if not check_disk_space():
166
- log_message("❌ Insufficient disk space even after cleanup")
167
- return False
168
-
169
- headers = {"Authorization": f"Bearer {HF_TOKEN}"}
170
- for attempt in range(max_retries):
171
- try:
172
- with requests.get(url, headers=headers, stream=True) as r:
173
- r.raise_for_status()
174
-
175
- # Check content length if available
176
- content_length = r.headers.get("content-length")
177
- if content_length:
178
- size_gb = int(content_length) / (1024**3)
179
- disk_info = get_disk_usage(".")
180
- if size_gb > disk_info["free"] - 0.5: # Leave 0.5GB buffer
181
- log_message(f'❌ File too large: {size_gb:.2f}GB, only {disk_info["free"]:.2f}GB free')
182
- return False
183
-
184
- with open(dest_path, "wb") as f:
185
- for chunk in r.iter_content(chunk_size=8192):
186
- f.write(chunk)
187
- return True
188
- except Exception as e:
189
- if attempt < max_retries - 1:
190
- time.sleep(2 ** attempt)
191
- continue
192
- log_message(f" Download failed after {max_retries} attempts: {e}")
193
- return False
194
- return False
195
-
196
- def is_multipart_rar(filename: str) -> bool:
197
- """Check if this is a multi-part RAR file"""
198
- return ".part" in filename.lower() and filename.lower().endswith(".rar")
199
-
200
- def get_rar_part_base(filename: str) -> str:
201
- """Get the base name for multi-part RAR files"""
202
- if ".part" in filename.lower():
203
- return filename.split(".part")[0]
204
- return filename.replace(".rar", "")
205
-
206
- def extract_with_retry(rar_path: str, output_dir: str, max_retries: int = 2) -> bool:
207
- """Extract RAR with retry and recovery, handling multi-part archives"""
208
- filename = os.path.basename(rar_path)
209
-
210
- # For multi-part RARs, we need the first part
211
- if is_multipart_rar(filename):
212
- base_name = get_rar_part_base(filename)
213
- first_part = f"{base_name}.part01.rar"
214
- first_part_path = os.path.join(os.path.dirname(rar_path), first_part)
215
-
216
- if not os.path.exists(first_part_path):
217
- log_message(f"⚠️ Multi-part RAR detected but first part not found: {first_part}")
218
- return False
219
-
220
- rar_path = first_part_path
221
- log_message(f"📦 Processing multi-part RAR starting with: {first_part}")
222
-
223
- for attempt in range(max_retries):
224
- try:
225
- # Test RAR first
226
- test_cmd = ["unrar", "t", rar_path]
227
- test_result = subprocess.run(test_cmd, capture_output=True, text=True)
228
- if test_result.returncode != 0:
229
- log_message(f"⚠️ RAR test failed: {test_result.stderr}")
230
- if attempt == max_retries - 1:
231
- return False
232
- continue
233
-
234
- # Extract RAR
235
- cmd = ["unrar", "x", "-o+", rar_path, output_dir]
236
- if attempt > 0: # Try recovery on subsequent attempts
237
- cmd.insert(2, "-kb")
238
-
239
- result = subprocess.run(cmd, capture_output=True, text=True)
240
- if result.returncode == 0:
241
- log_message(f"✅ Successfully extracted: {os.path.basename(rar_path)}")
242
- return True
243
- else:
244
- error_msg = result.stderr or result.stdout
245
- log_message(f"⚠️ Extraction attempt {attempt + 1} failed: {error_msg}")
246
-
247
- if "checksum error" in error_msg.lower() or "CRC failed" in error_msg:
248
- log_message(f"⚠️ Data corruption detected, attempt {attempt + 1}")
249
- elif result.returncode == 10:
250
- log_message(f"⚠️ No files to extract (exit code 10)")
251
- return False
252
- elif result.returncode == 1:
253
- log_message(f"⚠️ Non-fatal error (exit code 1)")
254
-
255
- except Exception as e:
256
- log_message(f"❌ Extraction exception: {str(e)}")
257
- if attempt == max_retries - 1:
258
- return False
259
- time.sleep(1)
260
-
261
- return False
262
-
263
- def ensure_dir(path):
264
- os.makedirs(path, exist_ok=True)
265
-
266
- def extract_frames(video_path, output_dir, fps=DEFAULT_FPS):
267
- """Extract frames from video at the specified frames per second (fps)."""
268
- log_message(f"[INFO] Extracting frames from {video_path} to {output_dir} at {fps} fps...")
269
- ensure_dir(output_dir)
270
- cap = cv2.VideoCapture(str(video_path))
271
- if not cap.isOpened():
272
- log_message(f"[ERROR] Failed to open video file: {video_path}")
273
- return 0
274
- video_fps = cap.get(cv2.CAP_PROP_FPS)
275
- if not video_fps or video_fps <= 0:
276
- video_fps = 30 # fallback if FPS is not available
277
- log_message(f"[WARN] Using fallback FPS: {video_fps}")
278
- frame_interval = int(round(video_fps / fps))
279
- frame_idx = 0
280
- saved_idx = 1
281
- total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
282
- log_message(f"[DEBUG] Total frames in video: {total_frames}")
283
- while cap.isOpened():
284
- ret, frame = cap.read()
285
- if not ret:
286
- break
287
- if frame_idx % frame_interval == 0:
288
- frame_name = f"{saved_idx:04d}.png"
289
- cv2.imwrite(str(Path(output_dir) / frame_name), frame)
290
- saved_idx += 1
291
- frame_idx += 1
292
- cap.release()
293
- log_message(f"Extracted {saved_idx-1} frames from {video_path} to {output_dir}")
294
- return saved_idx - 1
295
-
296
- def analyze_frame_with_florence2(image_path: str, prompt: str = "<CAPTION>") -> Dict:
297
- """Analyze a single frame using Florence-2 vision model."""
298
- if not vision_language_model_large or not vision_language_processor_large:
299
- return {
300
- "image": os.path.basename(image_path),
301
- "description": "[ERROR] Vision model not loaded."
302
- }
303
-
304
- image = Image.open(image_path).convert("RGB")
305
- inputs = vision_language_processor_large(images=image, text=prompt, return_tensors="pt").to(device)
306
-
307
- with torch.no_grad():
308
- generated_ids = vision_language_model_large.generate(
309
- input_ids=inputs["input_ids"],
310
- pixel_values=inputs["pixel_values"],
311
- max_new_tokens=512,
312
- do_sample=False,
313
- num_beams=3
314
- )
315
-
316
- generated_text = vision_language_processor_large.batch_decode(generated_ids, skip_special_tokens=False)[0]
317
- description = vision_language_processor_large.post_process_generation(
318
- generated_text,
319
- task="<CAPTION>",
320
- image_size=(image.width, image.height)
321
- )["<CAPTION>"]
322
-
323
- return {
324
- "image": os.path.basename(image_path),
325
- "description": description
326
- }
327
-
328
- def summarize_activities(frame_analyses: List[Dict]) -> Dict:
329
- """Summarize activities from frame analyses."""
330
- return {
331
- # "steps": [
332
- # {
333
- # "action": "Open Blender software",
334
- # "description": "User launches Blender 3D modeling application on their computer"
335
- # },
336
- # {
337
- # "action": "Create 3D object",
338
- # "description": "User works with a default cube object in the 3D viewport"
339
- # },
340
- # {
341
- # "action": "Manipulate 3D model",
342
- # "description": "User rotates and transforms the cube using mouse interactions"
343
- # },
344
- # {
345
- # "action": "Navigate interface",
346
- # "description": "User explores different tools and panels in the Blender interface"
347
- # }
348
- # ],
349
- # "high_level_goal": "Learning basic 3D modeling operations in Blender software",
350
- # "creative_actions": "3D object manipulation, interface navigation, basic modeling workflow",
351
- # "objects": ["computer", "monitor", "mouse", "keyboard", "Blender software", "3D cube", "desktop interface"],
352
- # "final_goal": "Introduction to Blender 3D modeling fundamentals and basic object manipulation"
353
- }
354
-
355
- def analyze_frames(frames_dir: str, output_json_path: str, prompt: Optional[str] = None) -> int:
356
- """Analyze all frames in directory using Florence-2 model."""
357
- log_message(f"[INFO] Analyzing frames in {frames_dir}...")
358
- frames_dir = Path(frames_dir).resolve()
359
- output_json_path = Path(output_json_path).resolve()
360
- ensure_dir(frames_dir)
361
- ensure_dir(output_json_path.parent)
362
-
363
- frame_analyses = []
364
- analyzed_count = 0
365
-
366
- for frame_file in sorted(frames_dir.glob("*.png")):
367
- analysis = analyze_frame_with_florence2(str(frame_file), prompt)
368
- frame_analyses.append(analysis)
369
- analyzed_count += 1
370
-
371
- # Generate summary
372
- summary = summarize_activities(frame_analyses)
373
-
374
- # Save results
375
- results = {
376
- "frame_analyses": frame_analyses,
377
- "summary": summary
378
- }
379
-
380
- try:
381
- with open(output_json_path, "w") as f:
382
- json.dump(results, f, indent=2)
383
- log_message(f"[SUCCESS] Analysis results saved to {output_json_path}")
384
- except Exception as e:
385
- log_message(f"[ERROR] Failed to write output JSON: {e}")
386
-
387
- return analyzed_count
388
-
389
-
390
- def process_rar_file(rar_path: str) -> bool:
391
- """Process a single RAR file - extract, then process videos for frames and vision analysis"""
392
- filename = os.path.basename(rar_path)
393
- processing_status["current_file"] = filename
394
-
395
- # Handle multi-part RAR naming
396
- if is_multipart_rar(filename):
397
- course_name = get_rar_part_base(filename)
398
- else:
399
- course_name = filename.replace(".rar", "")
400
-
401
- extract_dir = os.path.join(EXTRACT_FOLDER, course_name)
402
-
403
- try:
404
- log_message(f"🔄 Processing: {filename}")
405
-
406
- # Clean up any existing directory
407
- if os.path.exists(extract_dir):
408
- shutil.rmtree(extract_dir, ignore_errors=True)
409
-
410
- # Extract RAR
411
- os.makedirs(extract_dir, exist_ok=True)
412
- if not extract_with_retry(rar_path, extract_dir):
413
- raise Exception("RAR extraction failed")
414
-
415
- # Count extracted files
416
- file_count = 0
417
- video_files_found = []
418
- for root, dirs, files in os.walk(extract_dir):
419
- for file in files:
420
- file_count += 1
421
- if file.lower().endswith((".mp4", ".avi", ".mov", ".mkv")):
422
- video_files_found.append(os.path.join(root, file))
423
-
424
- processing_status["extracted_courses"] += 1
425
- log_message(f" Successfully extracted '{course_name}' ({file_count} files, {len(video_files_found)} videos)")
426
-
427
- # Process video files for frame extraction and vision analysis
428
- for video_path in video_files_found:
429
- video_filename = Path(video_path).name
430
- # Create a unique output directory for frames for each video
431
- frames_output_dir = os.path.join(FRAMES_OUTPUT_FOLDER, f"{course_name}_{video_filename.replace('.', '_')}_frames")
432
- ensure_dir(frames_output_dir)
433
-
434
- extracted_frames_count = extract_frames(video_path, frames_output_dir, fps=DEFAULT_FPS)
435
- processing_status["extracted_frames_count"] += extracted_frames_count
436
- if extracted_frames_count > 0:
437
- processing_status["extracted_videos"] += 1
438
- log_message(f"[INFO] Extracted {extracted_frames_count} frames from {video_filename}")
439
-
440
- # Perform vision analysis on the extracted frames
441
- analysis_output_json = os.path.join(ANALYSIS_OUTPUT_FOLDER, f"{course_name}_{video_filename.replace('.', '_')}_analysis.json")
442
- analyzed_frames = analyze_frames(frames_output_dir, analysis_output_json)
443
- processing_status["analyzed_frames_count"] += analyzed_frames
444
- log_message(f"[INFO] Analyzed {analyzed_frames} frames from {video_filename}")
445
- else:
446
- log_message(f"[WARN] No frames extracted from {video_filename}")
447
-
448
- return True
449
-
450
- except Exception as e:
451
- error_msg = str(e)
452
- log_message(f"❌ Processing failed: {error_msg}")
453
- log_failed_file(filename, error_msg)
454
- return False
455
-
456
- finally:
457
- processing_status["current_file"] = None
458
-
459
- def main_processing_loop(start_index: int = 0):
460
- """Main processing workflow - extraction, frame extraction, and vision analysis"""
461
- processing_status["is_running"] = True
462
-
463
- try:
464
- # Load state
465
- processed_rars = load_json_state(PROCESS_STATE_FILE, {"processed_rars": []})["processed_rars"]
466
- download_state = load_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": 0})
467
-
468
- # Use start_index if provided, otherwise use the saved state
469
- next_index = start_index if start_index > 0 else download_state["next_download_index"]
470
-
471
- log_message(f"📊 Starting from index {next_index}")
472
- log_message(f"📊 Previously processed: {len(processed_rars)} files")
473
-
474
- # Get file list
475
- try:
476
- files = list(hf_api.list_repo_files(repo_id=SOURCE_REPO_ID, repo_type="dataset"))
477
- rar_files = sorted([f for f in files if f.endswith(".rar")])
478
-
479
- processing_status["total_files"] = len(rar_files)
480
- log_message(f"📁 Found {len(rar_files)} RAR files in repository")
481
-
482
- if next_index >= len(rar_files):
483
- log_message("✅ All files have been processed!")
484
- return
485
-
486
- except Exception as e:
487
- log_message(f"❌ Failed to get file list: {str(e)}")
488
- return
489
-
490
- # Process only one file per run
491
- if next_index < len(rar_files):
492
- rar_file = rar_files[next_index]
493
- filename = os.path.basename(rar_file)
494
-
495
- if filename in processed_rars:
496
- log_message(f"⏭️ Skipping already processed: {filename}")
497
- processing_status["processed_files"] += 1
498
- # Move to next file
499
- next_index += 1
500
- save_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": next_index})
501
- log_message(f"📊 Moving to next file. Progress: {next_index}/{len(rar_files)}")
502
- return
503
-
504
- log_message(f"📥 Downloading: {filename}")
505
- dest_path = os.path.join(DOWNLOAD_FOLDER, filename)
506
-
507
- # Download file
508
- download_url = f"https://huggingface.co/datasets/{SOURCE_REPO_ID}/resolve/main/{rar_file}"
509
- if download_with_retry(download_url, dest_path):
510
- # Process file
511
- if process_rar_file(dest_path):
512
- processed_rars.append(filename)
513
- save_json_state(PROCESS_STATE_FILE, {"processed_rars": processed_rars})
514
- log_message(f"✅ Successfully processed: {filename}")
515
- processing_status["processed_files"] += 1
516
- else:
517
- log_message(f"❌ Failed to process: {filename}")
518
- processing_status["failed_files"] += 1
519
-
520
- # Clean up downloaded file
521
- try:
522
- os.remove(dest_path)
523
- log_message(f"🗑️ Cleaned up download: {filename}")
524
- except:
525
- pass
526
- else:
527
- log_message(f"❌ Failed to download: {filename}")
528
- processing_status["failed_files"] += 1
529
-
530
- # Update download state for next run
531
- next_index += 1
532
- save_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": next_index})
533
-
534
- # Status update
535
- log_message(f"📊 Progress: {next_index}/{len(rar_files)} files processed")
536
- log_message(f'📊 Extracted: {processing_status["extracted_courses"]} courses')
537
- log_message(f'📊 Videos Processed: {processing_status["extracted_videos"]}')
538
- log_message(f'📊 Frames Extracted: {processing_status["extracted_frames_count"]}')
539
- log_message(f'📊 Frames Analyzed: {processing_status["analyzed_frames_count"]}')
540
- log_message(f'📊 Failed: {processing_status["failed_files"]} files')
541
-
542
- if next_index < len(rar_files):
543
- log_message(f"🔄 Run the script again to process the next file: {os.path.basename(rar_files[next_index])}")
544
- else:
545
- log_message("🎉 All files have been processed!")
546
- else:
547
- log_message("✅ All files have been processed!")
548
-
549
- log_message("🎉 Processing complete!")
550
- log_message(f'📊 Final stats: {processing_status["extracted_courses"]} courses extracted, {processing_status["extracted_videos"]} videos processed, {processing_status["extracted_frames_count"]} frames extracted, {processing_status["analyzed_frames_count"]} frames analyzed')
551
-
552
- except KeyboardInterrupt:
553
- log_message("⏹️ Processing interrupted by user")
554
- except Exception as e:
555
- log_message(f"❌ Fatal error: {str(e)}")
556
- finally:
557
- processing_status["is_running"] = False
558
- cleanup_temp_files()
559
-
560
- # FastAPI Endpoints
561
- @app.post("/analyze-video")
562
- async def analyze_video_endpoint(
563
- file: UploadFile = File(...),
564
- fps: int = Form(DEFAULT_FPS),
565
- prompt: Optional[str] = Form(None)
566
- ):
567
- """Analyze a single video file and return frame-by-frame analysis."""
568
- if not file.filename.lower().endswith((".mp4", ".avi", ".mov", ".mkv")):
569
- return JSONResponse(status_code=400, content={
570
- "error": "File type not allowed",
571
- "allowed_types": [".mp4", ".avi", ".mov", ".mkv"]
572
- })
573
-
574
- with tempfile.TemporaryDirectory() as temp_dir:
575
- temp_dir_path = Path(temp_dir)
576
- file_path = temp_dir_path / file.filename
577
-
578
- with open(file_path, "wb") as buffer:
579
- shutil.copyfileobj(file.file, buffer)
580
-
581
- frames_dir = temp_dir_path / "frames"
582
- frame_count = extract_frames(file_path, frames_dir, fps)
583
-
584
- frame_analyses = []
585
- for frame_file in sorted(frames_dir.glob("*.png")):
586
- analysis = analyze_frame_with_florence2(str(frame_file), prompt)
587
- frame_analyses.append(analysis)
588
-
589
- summary = summarize_activities(frame_analyses)
590
-
591
- return JSONResponse(content={
592
- "video_filename": file.filename,
593
- "frame_count": frame_count,
594
- "fps": fps,
595
- "frame_analyses": frame_analyses,
596
- "summary": summary
597
- })
598
-
599
- @app.post("/analyze-archive")
600
- async def analyze_archive_endpoint(
601
- file: UploadFile = File(...),
602
- fps: int = Form(DEFAULT_FPS),
603
- prompt: Optional[str] = Form(None)
604
- ):
605
- """Analyze videos from RAR/ZIP archive and return frame-by-frame analysis."""
606
- if not file.filename.lower().endswith((".rar", ".zip")):
607
- return JSONResponse(status_code=400, content={
608
- "error": "File type not allowed",
609
- "allowed_types": [".rar", ".zip"]
610
- })
611
-
612
- with tempfile.TemporaryDirectory() as temp_dir:
613
- temp_dir_path = Path(temp_dir)
614
- file_path = temp_dir_path / file.filename
615
-
616
- with open(file_path, "wb") as buffer:
617
- shutil.copyfileobj(file.file, buffer)
618
-
619
- extract_dir = temp_dir_path / "extracted"
620
- video_files = []
621
-
622
- if file.filename.lower().endswith(".rar"):
623
- with rarfile.RarFile(file_path) as rf:
624
- rf.extractall(extract_dir)
625
- else:
626
- with zipfile.ZipFile(file_path) as zf:
627
- zf.extractall(extract_dir)
628
-
629
- # Find video files in extracted content
630
- for root, dirs, files in os.walk(extract_dir):
631
- for file in files:
632
- if file.lower().endswith((".mp4", ".avi", ".mov", ".mkv")):
633
- video_files.append(Path(root) / file)
634
-
635
- if not video_files:
636
- return JSONResponse(status_code=400, content={
637
- "error": "No video files found in archive"
638
- })
639
-
640
- results = []
641
- for video_path in video_files:
642
- video_name = video_path.name
643
- frames_dir = temp_dir_path / f"frames_{video_name}"
644
- frame_count = extract_frames(video_path, frames_dir, fps)
645
-
646
- frame_analyses = []
647
- for frame_file in sorted(frames_dir.glob("*.png")):
648
- analysis = analyze_frame_with_florence2(str(frame_file), prompt)
649
- frame_analyses.append(analysis)
650
-
651
- summary = summarize_activities(frame_analyses)
652
-
653
- results.append({
654
- "video_filename": video_name,
655
- "frame_count": frame_count,
656
- "fps": fps,
657
- "frame_analyses": frame_analyses,
658
- "summary": summary
659
- })
660
-
661
- return JSONResponse(content={
662
- "archive_filename": file.filename,
663
- "videos_processed": len(video_files),
664
- "results": results
665
- })
666
-
667
- @app.get("/health")
668
- async def health_check():
669
- """Health check endpoint."""
670
- return JSONResponse(content={
671
- "status": "healthy",
672
- "model": "Florence-2 (Mock)",
673
- "note": "Florence-2 model is mocked due to sandbox memory limitations."
674
- })
675
-
676
- @app.get("/status")
677
- async def get_processing_status():
678
- """Get current processing status."""
679
- return JSONResponse(content=processing_status)
680
-
681
- # Expose necessary functions and variables
682
- __all__ = [
683
- "main_processing_loop",
684
- "processing_status",
685
- "ANALYSIS_OUTPUT_FOLDER",
686
- "log_message",
687
- "send_email_with_attachment",
688
- "analyze_frames",
689
- "extract_frames",
690
- "DEFAULT_FPS",
691
- "ensure_dir"
692
- ]
 
1
+ import os
2
+ import json
3
+ import requests
4
+ import subprocess
5
+ import shutil
6
+ import time
7
+ import re
8
+ import threading
9
+ from typing import Dict, List, Set, Optional
10
+ from huggingface_hub import HfApi, list_repo_files
11
+
12
+ import cv2
13
+ import numpy as np
14
+ from pathlib import Path
15
+ import smtplib
16
+ from email.message import EmailMessage
17
+
18
+ # ==== CONFIGURATION ====
19
+ HF_TOKEN = os.getenv("HF_TOKEN", "")
20
+ SOURCE_REPO_ID = os.getenv("SOURCE_REPO", "Fred808/BG1")
21
+
22
+ # Path Configuration
23
+ DOWNLOAD_FOLDER = "downloads"
24
+ EXTRACT_FOLDER = "extracted"
25
+ FRAMES_OUTPUT_FOLDER = "extracted_frames"
26
+
27
+ os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)
28
+ os.makedirs(EXTRACT_FOLDER, exist_ok=True)
29
+ os.makedirs(FRAMES_OUTPUT_FOLDER, exist_ok=True)
30
+
31
+ # State Files
32
+ DOWNLOAD_STATE_FILE = "download_progress.json"
33
+ PROCESS_STATE_FILE = "process_progress.json"
34
+ FAILED_FILES_LOG = "failed_files.log"
35
+
36
+ # Processing Parameters
37
+ CHUNK_SIZE = 1
38
+ PROCESSING_DELAY = 2
39
+ MAX_RETRIES = 3
40
+ MIN_FREE_SPACE_GB = 2 # Minimum free space in GB before processing
41
+
42
+ # Frame Extraction Parameters
43
+ DEFAULT_FPS = 3 # Default frames per second for extraction
44
+
45
+ # Cursor Tracking Parameters
46
+
47
+
48
+ # Initialize HF API
49
+ hf_api = HfApi(token=HF_TOKEN)
50
+
51
+ # Global State
52
+ processing_status = {
53
+ "is_running": False,
54
+ "current_file": None,
55
+ "total_files": 0,
56
+ "processed_files": 0,
57
+ "failed_files": 0,
58
+ "extracted_courses": 0,
59
+ "extracted_videos": 0,
60
+ "last_update": None,
61
+ "logs": []
62
+ }
63
+
64
+ def log_message(message: str):
65
+ """Log messages with timestamp"""
66
+ timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
67
+ log_entry = f"[{timestamp}] {message}"
68
+ print(log_entry)
69
+ processing_status["logs"].append(log_entry)
70
+ processing_status["last_update"] = timestamp
71
+ if len(processing_status["logs"]) > 100:
72
+ processing_status["logs"] = processing_status["logs"][-100:]
73
+
74
+ def log_failed_file(filename: str, error: str):
75
+ """Log failed files to persistent file"""
76
+ with open(FAILED_FILES_LOG, "a") as f:
77
+ f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {filename}: {error}\n")
78
+
79
+ def get_disk_usage(path: str) -> Dict[str, float]:
80
+ """Get disk usage statistics in GB"""
81
+ statvfs = os.statvfs(path)
82
+ total = statvfs.f_frsize * statvfs.f_blocks / (1024**3)
83
+ free = statvfs.f_frsize * statvfs.f_bavail / (1024**3)
84
+ used = total - free
85
+ return {"total": total, "free": free, "used": used}
86
+
87
+ def check_disk_space(path: str = ".") -> bool:
88
+ """Check if there's enough disk space"""
89
+ disk_info = get_disk_usage(path)
90
+ if disk_info["free"] < MIN_FREE_SPACE_GB:
91
+ log_message(f'⚠️ Low disk space: {disk_info["free"]:.2f}GB free, {disk_info["used"]:.2f}GB used')
92
+ return False
93
+ return True
94
+
95
+ def cleanup_temp_files():
96
+ """Clean up temporary files to free space"""
97
+ log_message("🧹 Cleaning up temporary files...")
98
+
99
+ # Clean old downloads (keep only current processing file)
100
+ current_file = processing_status.get("current_file")
101
+ for file in os.listdir(DOWNLOAD_FOLDER):
102
+ if file != current_file and file.endswith((".rar", ".zip")):
103
+ try:
104
+ os.remove(os.path.join(DOWNLOAD_FOLDER, file))
105
+ log_message(f"🗑️ Removed old download: {file}")
106
+ except:
107
+ pass
108
+
109
+ def load_json_state(file_path: str, default_value):
110
+ """Load state from JSON file"""
111
+ if os.path.exists(file_path):
112
+ try:
113
+ with open(file_path, "r") as f:
114
+ return json.load(f)
115
+ except json.JSONDecodeError:
116
+ log_message(f"⚠️ Corrupted state file: {file_path}")
117
+ return default_value
118
+
119
+ def save_json_state(file_path: str, data):
120
+ """Save state to JSON file"""
121
+ with open(file_path, "w") as f:
122
+ json.dump(data, f, indent=2)
123
+
124
+ def download_with_retry(url: str, dest_path: str, max_retries: int = 3) -> bool:
125
+ """Download file with retry logic and disk space checking"""
126
+ if not check_disk_space():
127
+ cleanup_temp_files()
128
+ if not check_disk_space():
129
+ log_message("❌ Insufficient disk space even after cleanup")
130
+ return False
131
+
132
+ headers = {"Authorization": f"Bearer {HF_TOKEN}"}
133
+ for attempt in range(max_retries):
134
+ try:
135
+ with requests.get(url, headers=headers, stream=True) as r:
136
+ r.raise_for_status()
137
+
138
+ # Check content length if available
139
+ content_length = r.headers.get("content-length")
140
+ if content_length:
141
+ size_gb = int(content_length) / (1024**3)
142
+ disk_info = get_disk_usage(".")
143
+ if size_gb > disk_info["free"] - 0.5: # Leave 0.5GB buffer
144
+ log_message(f'❌ File too large: {size_gb:.2f}GB, only {disk_info["free"]:.2f}GB free')
145
+ return False
146
+
147
+ with open(dest_path, "wb") as f:
148
+ for chunk in r.iter_content(chunk_size=8192):
149
+ f.write(chunk)
150
+ return True
151
+ except Exception as e:
152
+ if attempt < max_retries - 1:
153
+ time.sleep(2 ** attempt)
154
+ continue
155
+ log_message(f"❌ Download failed after {max_retries} attempts: {e}")
156
+ return False
157
+ return False
158
+
159
+ def is_multipart_rar(filename: str) -> bool:
160
+ """Check if this is a multi-part RAR file"""
161
+ return ".part" in filename.lower() and filename.lower().endswith(".rar")
162
+
163
+ def get_rar_part_base(filename: str) -> str:
164
+ """Get the base name for multi-part RAR files"""
165
+ if ".part" in filename.lower():
166
+ return filename.split(".part")[0]
167
+ return filename.replace(".rar", "")
168
+
169
+ def extract_with_retry(rar_path: str, output_dir: str, max_retries: int = 2) -> bool:
170
+ """Extract RAR with retry and recovery, handling multi-part archives"""
171
+ filename = os.path.basename(rar_path)
172
+
173
+ # For multi-part RARs, we need the first part
174
+ if is_multipart_rar(filename):
175
+ base_name = get_rar_part_base(filename)
176
+ first_part = f"{base_name}.part01.rar"
177
+ first_part_path = os.path.join(os.path.dirname(rar_path), first_part)
178
+
179
+ if not os.path.exists(first_part_path):
180
+ log_message(f"⚠️ Multi-part RAR detected but first part not found: {first_part}")
181
+ return False
182
+
183
+ rar_path = first_part_path
184
+ log_message(f"📦 Processing multi-part RAR starting with: {first_part}")
185
+
186
+ for attempt in range(max_retries):
187
+ try:
188
+ # Test RAR first
189
+ test_cmd = ["unrar", "t", rar_path]
190
+ test_result = subprocess.run(test_cmd, capture_output=True, text=True)
191
+ if test_result.returncode != 0:
192
+ log_message(f"⚠️ RAR test failed: {test_result.stderr}")
193
+ if attempt == max_retries - 1:
194
+ return False
195
+ continue
196
+
197
+ # Extract RAR
198
+ cmd = ["unrar", "x", "-o+", rar_path, output_dir]
199
+ if attempt > 0: # Try recovery on subsequent attempts
200
+ cmd.insert(2, "-kb")
201
+
202
+ result = subprocess.run(cmd, capture_output=True, text=True)
203
+ if result.returncode == 0:
204
+ log_message(f"✅ Successfully extracted: {os.path.basename(rar_path)}")
205
+ return True
206
+ else:
207
+ error_msg = result.stderr or result.stdout
208
+ log_message(f"⚠️ Extraction attempt {attempt + 1} failed: {error_msg}")
209
+
210
+ if "checksum error" in error_msg.lower() or "CRC failed" in error_msg:
211
+ log_message(f"⚠️ Data corruption detected, attempt {attempt + 1}")
212
+ elif result.returncode == 10:
213
+ log_message(f"⚠️ No files to extract (exit code 10)")
214
+ return False
215
+ elif result.returncode == 1:
216
+ log_message(f"⚠️ Non-fatal error (exit code 1)")
217
+
218
+ except Exception as e:
219
+ log_message(f"❌ Extraction exception: {str(e)}")
220
+ if attempt == max_retries - 1:
221
+ return False
222
+ time.sleep(1)
223
+
224
+ return False
225
+
226
+ # --- Frame Extraction Utilities ---
227
+ def ensure_dir(path):
228
+ os.makedirs(path, exist_ok=True)
229
+
230
+ def extract_frames(video_path, output_dir, fps=DEFAULT_FPS):
231
+ """Extract frames from video at the specified frames per second (fps)."""
232
+ log_message(f"[INFO] Extracting frames from {video_path} to {output_dir} at {fps} fps...")
233
+ ensure_dir(output_dir)
234
+ cap = cv2.VideoCapture(str(video_path))
235
+ if not cap.isOpened():
236
+ log_message(f"[ERROR] Failed to open video file: {video_path}")
237
+ return 0
238
+ video_fps = cap.get(cv2.CAP_PROP_FPS)
239
+ # log_message(f"[DEBUG] Video FPS: {video_fps}")
240
+ if not video_fps or video_fps <= 0:
241
+ video_fps = 30 # fallback if FPS is not available
242
+ log_message(f"[WARN] Using fallback FPS: {video_fps}")
243
+ frame_interval = int(round(video_fps / fps))
244
+ # log_message(f"[DEBUG] Frame interval: {frame_interval}")
245
+ frame_idx = 0
246
+ saved_idx = 1
247
+ total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
248
+ log_message(f"[DEBUG] Total frames in video: {total_frames}")
249
+ while cap.isOpened():
250
+ ret, frame = cap.read()
251
+ if not ret:
252
+ # log_message(f"[DEBUG] No more frames to read at frame_idx {frame_idx}.")
253
+ break
254
+ if frame_idx % frame_interval == 0:
255
+ frame_name = f"{saved_idx:04d}.png"
256
+ cv2.imwrite(str(Path(output_dir) / frame_name), frame)
257
+ # log_message(f"[DEBUG] Saved frame {frame_idx} as {frame_name}")
258
+ saved_idx += 1
259
+ frame_idx += 1
260
+ cap.release()
261
+ log_message(f"Extracted {saved_idx-1} frames from {video_path} to {output_dir}")
262
+ return saved_idx - 1
263
+
264
+
265
+
266
+
267
+
268
+ def process_rar_file(rar_path: str) -> bool:
269
+ """Process a single RAR file - extract, then process videos for frames"""
270
+ filename = os.path.basename(rar_path)
271
+ processing_status["current_file"] = filename
272
+
273
+ # Handle multi-part RAR naming
274
+ if is_multipart_rar(filename):
275
+ course_name = get_rar_part_base(filename)
276
+ else:
277
+ course_name = filename.replace(".rar", "")
278
+
279
+ extract_dir = os.path.join(EXTRACT_FOLDER, course_name)
280
+
281
+ try:
282
+ log_message(f"🔄 Processing: {filename}")
283
+
284
+ # Clean up any existing directory
285
+ if os.path.exists(extract_dir):
286
+ shutil.rmtree(extract_dir, ignore_errors=True)
287
+
288
+ # Extract RAR
289
+ os.makedirs(extract_dir, exist_ok=True)
290
+ if not extract_with_retry(rar_path, extract_dir):
291
+ raise Exception("RAR extraction failed")
292
+
293
+ # Count extracted files
294
+ file_count = 0
295
+ video_files_found = []
296
+ for root, dirs, files in os.walk(extract_dir):
297
+ for file in files:
298
+ file_count += 1
299
+ if file.lower().endswith((".mp4", ".avi", ".mov", ".mkv")):
300
+ video_files_found.append(os.path.join(root, file))
301
+
302
+ processing_status["extracted_courses"] += 1
303
+ log_message(f"✅ Successfully extracted '{course_name}' ({file_count} files, {len(video_files_found)} videos)")
304
+
305
+ # Process video files for frame extraction
306
+ for video_path in video_files_found:
307
+ video_filename = Path(video_path).name
308
+ # Unique output directory for frames
309
+ frames_output_dir = os.path.join(
310
+ FRAMES_OUTPUT_FOLDER,
311
+ f"{course_name}_{video_filename.replace('.', '_')}_frames"
312
+ )
313
+ ensure_dir(frames_output_dir)
314
+
315
+ # 🔥 Extract frames here
316
+ frame_count = extract_frames(video_path, frames_output_dir, fps=DEFAULT_FPS)
317
+ processing_status["extracted_videos"] += 1
318
+
319
+ if frame_count == 0:
320
+ log_message(f"⚠️ No frames extracted from {video_filename}")
321
+ else:
322
+ log_message(f"✅ {frame_count} frames extracted from {video_filename}")
323
+
324
+ return True
325
+
326
+ except Exception as e:
327
+ error_msg = str(e)
328
+ log_message(f"❌ Processing failed: {error_msg}")
329
+ log_failed_file(filename, error_msg)
330
+ return False
331
+
332
+ finally:
333
+ processing_status["current_file"] = None
334
+
335
+
336
+ def main_processing_loop(start_index: int = 0):
337
+ """Main processing workflow - extraction, frame extraction, and cursor tracking"""
338
+ processing_status["is_running"] = True
339
+
340
+ try:
341
+ # Load state
342
+ processed_rars = load_json_state(PROCESS_STATE_FILE, {"processed_rars": []})["processed_rars"]
343
+ download_state = load_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": 4})
344
+
345
+ # Use start_index if provided, otherwise use the saved state
346
+ next_index = start_index if start_index > 0 else download_state["next_download_index"]
347
+
348
+ log_message(f"📊 Starting from index {next_index}")
349
+ log_message(f"📊 Previously processed: {len(processed_rars)} files")
350
+
351
+ # Get file list
352
+ try:
353
+ files = list(hf_api.list_repo_files(repo_id=SOURCE_REPO_ID, repo_type="dataset"))
354
+ rar_files = sorted([f for f in files if f.endswith(".rar")])
355
+
356
+ processing_status["total_files"] = len(rar_files)
357
+ log_message(f"📁 Found {len(rar_files)} RAR files in repository")
358
+
359
+ if next_index >= len(rar_files):
360
+ log_message("✅ All files have been processed!")
361
+ return
362
+
363
+ except Exception as e:
364
+ log_message(f"❌ Failed to get file list: {str(e)}")
365
+ return
366
+
367
+ # Process only one file per run
368
+ if next_index < len(rar_files):
369
+ rar_file = rar_files[next_index]
370
+ filename = os.path.basename(rar_file)
371
+
372
+ if filename in processed_rars:
373
+ log_message(f"⏭️ Skipping already processed: {filename}")
374
+ processing_status["processed_files"] += 1
375
+ # Move to next file
376
+ next_index += 1
377
+ save_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": next_index})
378
+ log_message(f"📊 Moving to next file. Progress: {next_index}/{len(rar_files)}")
379
+ return
380
+
381
+ log_message(f"📥 Downloading: {filename}")
382
+ dest_path = os.path.join(DOWNLOAD_FOLDER, filename)
383
+
384
+ # Download file
385
+ download_url = f"https://huggingface.co/datasets/{SOURCE_REPO_ID}/resolve/main/{rar_file}"
386
+ if download_with_retry(download_url, dest_path):
387
+ # Process file
388
+ if process_rar_file(dest_path):
389
+ processed_rars.append(filename)
390
+ save_json_state(PROCESS_STATE_FILE, {"processed_rars": processed_rars})
391
+ log_message(f" Successfully processed: {filename}")
392
+ processing_status["processed_files"] += 1
393
+ else:
394
+ log_message(f"❌ Failed to process: {filename}")
395
+ processing_status["failed_files"] += 1
396
+
397
+ # Clean up downloaded file
398
+ try:
399
+ os.remove(dest_path)
400
+ log_message(f"🗑️ Cleaned up download: {filename}")
401
+ except:
402
+ pass
403
+ else:
404
+ log_message(f" Failed to download: {filename}")
405
+ processing_status["failed_files"] += 1
406
+
407
+ # Update download state for next run
408
+ next_index += 1
409
+ save_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": next_index})
410
+
411
+ # Status update log_message(f"📊 Frames Extracted: {processing_status["extracted_frames_count"]}")
412
+ if next_index < len(rar_files):
413
+ log_message(f"🔄 Run the script again to process the next file: {os.path.basename(rar_files[next_index])}")
414
+ else:
415
+ log_message("🎉 All files have been processed!")
416
+ else:
417
+ log_message("✅ All files have been processed!")
418
+
419
+ log_message("🎉 Processing complete!")
420
+ log_message(f'📊 Final stats: {processing_status["extracted_courses"]} courses extracted, {processing_status["extracted_videos"]} videos processed, frames extracted')
421
+
422
+ except KeyboardInterrupt:
423
+ log_message("⏹️ Processing interrupted by user")
424
+ except Exception as e:
425
+ log_message(f" Fatal error: {str(e)}")
426
+ finally:
427
+ processing_status["is_running"] = False
428
+ cleanup_temp_files()
429
+
430
+ # Expose necessary functions and variables for download_api.py
431
+ __all__ = [
432
+ "main_processing_loop",
433
+ "processing_status",
434
+ "log_message",
435
+ "extract_frames",
436
+ "DEFAULT_FPS",
437
+ "ensure_dir"
438
+ ]
439
+