Update app.py

app.py

@@ -15,17 +15,23 @@ st.set_page_config(page_title="Solar Panel Fault Detection", layout="wide")
 
 # Title and description
 st.title("Solar Panel Fault Detection PoC")
-st.write("Upload a thermal video (MP4) of a solar panel to detect thermal, dust, and power generation faults.")
+st.write("Upload a thermal video (MP4) to detect thermal, dust, and power generation faults.")
+
+# 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)
+resize_enabled = st.sidebar.checkbox("Resize Frames (faster processing)", value=True)
+resize_width = 640 if resize_enabled else None
 
 # Load model and processor
 @st.cache_resource
 def load_model():
-    # Suppress warning about unused weights
     warnings.filterwarnings("ignore", message="Some weights of the model checkpoint.*were not used")
     logging.set_verbosity_error()
     
-    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+    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()
@@ -33,9 +39,19 @@ def load_model():
 
 processor, model, device = load_model()
 
+# Function to resize frame
+def resize_frame(frame, width=None):
+    if width is 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)
+
 # 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():
             outputs = model(**inputs)
@@ -71,13 +87,17 @@ 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()
+        
         return annotated_frames, all_faults
     except Exception as e:
         st.error(f"Error during fault detection: {str(e)}")
         return [], []
 
 # Function to process video
-async def process_video(video_path, frame_skip=5, batch_size=4):
+async def process_video(video_path, frame_skip, batch_size):
     try:
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
@@ -89,9 +109,13 @@ async def process_video(video_path, frame_skip=5, batch_size=4):
         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
+        
         output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
+        out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height))
         
         video_faults = {"Thermal Fault": False, "Dust Fault": False, "Power Generation Fault": False}
         frame_count = 0
@@ -104,9 +128,12 @@ async def process_video(video_path, frame_skip=5, batch_size=4):
             
             while cap.isOpened():
                 ret, frame = cap.read()
-                if not, break
+                if not ret:
+                    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
                     processed_frames += 1
@@ -166,7 +193,7 @@ 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=5, batch_size=4))
+        output_path, video_faults = loop.run_until_complete(process_video(tfile.name, frame_skip, batch_size))
         
         if output_path and video_faults:
             st.subheader("Fault Detection Results")