optimize

config.py

@@ -1,7 +1,8 @@
 # config.py
 import torch
-UPLOAD_DIR = "/code/static/uploads"
-OUTPUT_DIR = "/code/static/outputs"
+import os
+UPLOAD_DIR = os.path.join(os.getcwd(), "static/uploads")
+OUTPUT_DIR = os.path.join(os.getcwd(), "static/outputs")
 FRAME_RATE = 15
 SCORE_THRESHOLD = 0.4
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

routes/summarize.py

@@ -1,8 +1,7 @@
 from fastapi import APIRouter, UploadFile, File
 from fastapi.responses import JSONResponse
 from utils.file_utils import save_uploaded_file
-from services.extractor import extract_features
-from services.model_loader import load_model
+from services.extractor import extract_frames, extract_features
 from services.summarizer import get_scores, get_selected_indices, save_summary_video
 from config import UPLOAD_DIR, OUTPUT_DIR
 
@@ -13,15 +12,27 @@ def summarize_video(video: UploadFile = File(...)):
     if not video.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
         return JSONResponse(content={"error": "Unsupported file format"}, status_code=400)
 
+    print("\n-----------> Uploading Video ....")
     video_path = save_uploaded_file(video, UPLOAD_DIR)
-    features, picks = extract_features(video_path)
-    model = load_model("Model/epoch-199.pkl")
-    scores = get_scores(model, features)
+
+    print("\n-----------> Extracting Frames ....")
+    frames, picks = extract_frames(video_path)
+
+    print("\n-----------> Extracting Features ....")
+    features = extract_features(frames)
+
+    print("\n-----------> Getting Scores ....")
+    scores = get_scores(features)
+
+    print("\n-----------> Selecting Indices ....")
     selected = get_selected_indices(scores, picks)
     output_path = f"{OUTPUT_DIR}/summary_{video.filename}"
+
+    print("\n-----------> Saving Video ....")
     save_summary_video(video_path, selected, output_path)
     summary_url = f"/static/outputs/summary_{video.filename}"
 
+    print("\n-----------> Returning Response ....")
     return JSONResponse(content={
         "message": "Summarization complete",
         "summary_video_url": summary_url

services/extractor.py

@@ -4,6 +4,7 @@ import numpy as np
 from PIL import Image
 from torchvision import models, transforms
 from config import DEVICE, FRAME_RATE
+from tqdm import tqdm
 
 # Load GoogLeNet once
 from torchvision.models import GoogLeNet_Weights
@@ -40,23 +41,26 @@ transform = transforms.Compose([
     )
 ])
 
-def extract_features(video_path):
+def extract_frames(video_path):
     cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    picks, frames = [], []
-    count = 0
+    frames = []
+    indices = []
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    total_frames = 100 # TEMP
 
-    while cap.isOpened():
+    for idx in tqdm(range(0, total_frames, FRAME_RATE)):
+        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
         ret, frame = cap.read()
         if not ret:
             break
-        if int(count % round(fps // FRAME_RATE)) == 0:
-            image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-            input_tensor = transform(image).unsqueeze(0).to(DEVICE)
-            with torch.no_grad():
-                feature = feature_extractor(input_tensor).squeeze(0).cpu().numpy()
-                frames.append(feature)
-                picks.append(count)
-        count += 1
+        frames.append(Image.fromarray(frame))
+        indices.append(idx)
+
     cap.release()
-    return np.stack(frames), picks
+    return frames, indices
+
+def extract_features(frames):
+    features = [transform(frame) for frame in frames]
+    features = torch.stack(features).to(DEVICE)
+    features = feature_extractor(features)
+    return features

services/summarizer.py

@@ -1,16 +1,19 @@
 import cv2
 import torch
 from config import SCORE_THRESHOLD
+from services.model_loader import load_model
 
-def get_scores(model, features):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = load_model("Model/epoch-199.pkl")
+model = model.to(device)
+model = model.eval()
+
+def get_scores(features):
+    # features.shape: (N, 1024)
     with torch.no_grad():
-        features_tensor = torch.tensor(features, dtype=torch.float32).to(device)
-        scores, _ = model(features_tensor)
+        scores, _ = model(features)
     return scores.squeeze().cpu().numpy()
 
-
 def get_selected_indices(scores, picks, threshold=SCORE_THRESHOLD):
     return [picks[i] for i, score in enumerate(scores) if score >= threshold]