Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import DetrImageProcessor, DetrForObjectDetection
+from transformers import YolosImageProcessor, YolosForObjectDetection
 import cv2
 import numpy as np
 import tempfile
@@ -19,23 +19,35 @@ st.write("Upload a thermal video (MP4) to detect thermal, dust, and power genera
 
 # UI controls for optimization parameters
 st.sidebar.header("Analysis Settings")
-frame_skip = st.sidebar.slider("Frame Skip (higher = faster, less thorough)", min_value=1, max_value=20, value=10)
-batch_size = st.sidebar.slider("Batch Size (adjust for hardware)", min_value=1, max_value=16, value=8)
+frame_skip = st.sidebar.slider("Frame Skip (higher = faster, less thorough)", min_value=1, max_value=30, value=15)
+batch_size = st.sidebar.slider("Batch Size (adjust for hardware)", min_value=1, max_value=16, value=12)
 resize_enabled = st.sidebar.checkbox("Resize Frames (faster processing)", value=True)
 resize_width = 640 if resize_enabled else None
+quantize_model = st.sidebar.checkbox("Quantize Model (faster on CPU)", value=False)
 
 # Load model and processor
 @st.cache_resource
-def load_model():
+def load_model(quantize=quantize_model):
     warnings.filterwarnings("ignore", message="Some weights of the model checkpoint.*were not used")
     logging.set_verbosity_error()
     
-    processor = DetrImageProcessor.from_pretrained("hustvl/yolos-tiny")
-    model = DetrForObjectDetection.from_pretrained("hustvl/yolos-tiny")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    model.eval()
-    return processor, model, device
+    try:
+        processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny")
+        model = YolosForObjectDetection.from_pretrained("hustvl/yolos-tiny")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device)
+        
+        # Apply dynamic quantization for CPU if enabled
+        if quantize and device.type == "cpu":
+            model = torch.quantization.quantize_dynamic(
+                model, {torch.nn.Linear}, dtype=torch.qint8
+            )
+        
+        model.eval()
+        return processor, model, device
+    except Exception as e:
+        st.error(f"Failed to load model: {str(e)}. Please check your internet connection or clear the cache (~/.cache/huggingface/hub).")
+        raise
 
 processor, model, device = load_model()
 
@@ -50,7 +62,6 @@ def resize_frame(frame, width=None):
 # Function to process a batch of frames
 async def detect_faults_batch(frames, processor, model, device):
     try:
-        # Resize frames if enabled
         frames = [resize_frame(frame, resize_width) for frame in frames]
         inputs = processor(images=frames, return_tensors="pt").to(device)
         with torch.no_grad():
@@ -87,7 +98,6 @@ async def detect_faults_batch(frames, processor, model, device):
             annotated_frames.append(annotated_frame)
             all_faults.append(faults)
         
-        # Clear GPU memory
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         
@@ -109,7 +119,6 @@ async def process_video(video_path, frame_skip, batch_size):
         fps = int(cap.get(cv2.CAP_PROP_FPS))
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        # Adjust output size if resizing
         out_width = resize_width if resize_width else frame_width
         out_height = int(out_width * frame_height / frame_width) if resize_width else frame_height
         
@@ -132,7 +141,6 @@ async def process_video(video_path, frame_skip, batch_size):
                     break
                 
                 if frame_count % frame_skip != 0:
-                    # Resize frame for output if needed
                     frame = resize_frame(frame, resize_width)
                     out.write(frame)
                     frame_count += 1