Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import YolosImageProcessor, YolosForObjectDetection
+from transformers import DetrImageProcessor, DetrForObjectDetection
 import cv2
 import numpy as np
 import tempfile
@@ -19,11 +19,11 @@ 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=30, value=15)
-batch_size = st.sidebar.slider("Batch Size (adjust for hardware)", min_value=1, max_value=16, value=12)
+frame_skip = st.sidebar.slider("Frame Skip (higher = faster, less thorough)", min_value=1, max_value=50, value=30)
+batch_size = st.sidebar.slider("Batch Size (adjust for hardware)", min_value=1, max_value=32, value=16 if torch.cuda.is_available() else 8)
 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)
+resize_width = 512 if resize_enabled else None
+quantize_model = st.sidebar.checkbox("Quantize Model (faster, esp. on CPU)", value=True)
 
 # Load model and processor
 @st.cache_resource
@@ -32,12 +32,12 @@ def load_model(quantize=quantize_model):
     logging.set_verbosity_error()
     
     try:
-        processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny")
-        model = YolosForObjectDetection.from_pretrained("hustvl/yolos-tiny")
+        processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+        model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         model.to(device)
         
-        # Apply dynamic quantization for CPU if enabled
+        # Apply dynamic quantization if enabled
         if quantize and device.type == "cpu":
             model = torch.quantization.quantize_dynamic(
                 model, {torch.nn.Linear}, dtype=torch.qint8
@@ -46,7 +46,7 @@ def load_model(quantize=quantize_model):
         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).")
+        st.error(f"Failed to load model: {str(e)}. Check internet connection or cache (~/.cache/huggingface/hub).")
         raise
 
 processor, model, device = load_model()
@@ -57,7 +57,7 @@ def resize_frame(frame, width=None):
         return frame
     aspect_ratio = frame.shape[1] / frame.shape[0]
     height = int(width / aspect_ratio)
-    return cv2.resize(frame, (width, height), interpolation=cv2.INTER_AREA)
+    return cv2.resize(frame, (width, height), interpolation=cv2.INTER_LINEAR)
 
 # Function to process a batch of frames
 async def detect_faults_batch(frames, processor, model, device):
@@ -162,7 +162,7 @@ async def process_video(video_path, frame_skip, batch_size):
                     frames_batch = []
                     processed_frames += batch_size
                     progress.progress(min(processed_frames / total_frames, 1.0))
-                
+ GROUP BY 
                 frame_count += 1
             
             if frames_batch:
@@ -200,8 +200,13 @@ if uploaded_file is not None:
         
         st.video(tfile.name, format="video/mp4")
         
-        loop = asyncio.get_event_loop()
-        output_path, video_faults = loop.run_until_complete(process_video(tfile.name, frame_skip, batch_size))
+        # Create a new event loop for Streamlit's ScriptRunner thread
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            output_path, video_faults = loop.run_until_complete(process_video(tfile.name, frame_skip, batch_size))
+        finally:
+            loop.close()
         
         if output_path and video_faults:
             st.subheader("Fault Detection Results")