Update app.py

app.py

@@ -5,16 +5,13 @@ import numpy as np
 from ultralytics import YOLO
 from tqdm import tqdm
 
-# In Spaces, you can rely on ultralytics to fetch yolov8n.pt if not present locally.
+# Pre-load models from the repo folder
+extract_model = YOLO("best.pt")
+detect_model  = YOLO("yolov8n.pt")
 
-def process_video(video_path, "best.pt"):
-    # video_path and best_model_path are now plain file paths
+def process_video(video_path):
+    # Prepare output folder
     os.makedirs("frames", exist_ok=True)
-    # load models exactly as before...
-    extract_model = YOLO("best.pt")
-    detect_model  = YOLO("yolov8n.pt")
-    # ...
-
 
     # --- Step 1: Extract clean frames ---
     cap = cv2.VideoCapture(video_path)
@@ -78,33 +75,29 @@ def process_video(video_path, "best.pt"):
         masks.append(m)
 
     sum_img = np.zeros_like(aligned[0], dtype=np.float32)
-    count = np.zeros(aligned[0].shape[:2], dtype=np.float32)
+    count   = np.zeros(aligned[0].shape[:2], dtype=np.float32)
     for f, m in zip(aligned, masks):
-        inv = cv2.bitwise_not(m)
+        inv    = cv2.bitwise_not(m)
         masked = cv2.bitwise_and(f, f, mask=inv)
         sum_img += masked.astype(np.float32)
-        count += (inv>0).astype(np.float32)
+        count   += (inv>0).astype(np.float32)
 
     count[count==0] = 1
     selective = (sum_img / count[:,:,None]).astype(np.uint8)
     cv2.imwrite("fused_board_selective.jpg", selective)
 
     # --- Step 5: Sharpen final result ---
-    blur = cv2.GaussianBlur(selective, (5,5), 0)
+    blur  = cv2.GaussianBlur(selective, (5,5), 0)
     sharp = cv2.addWeighted(selective, 1.5, blur, -0.5, 0)
     cv2.imwrite("sharpened_board_color.jpg", sharp)
 
     return "clean_board.jpg", "fused_board_selective.jpg", "sharpened_board_color.jpg"
 
-
 demo = gr.Interface(
     fn=process_video,
     inputs=[
-        gr.Video(label="Upload Classroom Video (.mp4)", type="filepath"),
-        gr.File(
-            label="Upload fine-tuned YOLO model (best.pt)",
-            file_types=['.pt'],
-            file_count="single",
+        gr.Video(
+            label="Upload Classroom Video (.mp4)",
             type="filepath"
         )
     ],
@@ -115,13 +108,11 @@ demo = gr.Interface(
     ],
     title="📹 Classroom Board Cleaner",
     description=(
-        "1️⃣ Upload a video and your `best.pt` model\n"
-        "2️⃣ Process to extract, align, mask, fuse & sharpen\n"
-        "3️⃣ View three stages of output"
+        "1️⃣ Upload your classroom video (no model upload needed)\n"
+        "2️⃣ Automatic extraction, alignment, masking, fusion & sharpening\n"
+        "3️⃣ View three stages of the cleaned board output"
     )
 )
 
-
-
 if __name__ == "__main__":
     demo.launch()