Update app.py

app.py

@@ -375,26 +375,31 @@ def analyze_video_activities(video_path):
 @spaces.GPU
 def process_diary(day, date, total_people, total_machinery, machinery_types, activities, media_source, local_file, azure_blob):
     """Process the site diary entry with media from local file or Azure Blob Storage."""
-    if media_source == "Local File":
-        if local_file is None:
-            return [day, date, "No media uploaded", "No media uploaded", "No media uploaded", "No media uploaded", None]
-        media_path = local_file  # local_file is a string path in Gradio
-    else:  # Azure Blob
-        if not azure_blob:
-            return [day, date, "No blob selected", "No blob selected", "No blob selected", "No blob selected", None]
-        try:
-            blob_client = blob_service_client.get_blob_client(container=CONTAINER_NAME, blob=azure_blob)
-            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(azure_blob)[1]) as temp_file:
-                temp_path = temp_file.name
-                blob_data = blob_client.download_blob()
-                blob_data.readinto(temp_file)
-            media_path = temp_path
-        except Exception as e:
-            print(f"Error downloading blob: {str(e)}")
-            return [day, date, "Error downloading blob", "Error downloading blob", "Error downloading blob", "Error downloading blob", None]
-
-    # Process the media file
+    import torch
+    import gc
+
+    # Initialize media_path
+    media_path = None
     try:
+        if media_source == "Local File":
+            if local_file is None:
+                return [day, date, "No media uploaded", "No media uploaded", "No media uploaded", "No media uploaded", None]
+            media_path = local_file  # local_file is a string path in Gradio
+        else:  # Azure Blob
+            if not azure_blob:
+                return [day, date, "No blob selected", "No blob selected", "No blob selected", "No blob selected", None]
+            try:
+                blob_client = blob_service_client.get_blob_client(container=CONTAINER_NAME, blob=azure_blob)
+                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(azure_blob)[1]) as temp_file:
+                    temp_path = temp_file.name
+                    blob_data = blob_client.download_blob()
+                    blob_data.readinto(temp_file)
+                media_path = temp_path
+            except Exception as e:
+                print(f"Error downloading blob: {str(e)}")
+                return [day, date, "Error downloading blob", "Error downloading blob", "Error downloading blob", "Error downloading blob", None]
+
+        # Process the media file
         file_ext = get_file_extension(media_path)
         if not (is_image(media_path) or is_video(media_path)):
             raise ValueError(f"Unsupported file type: {file_ext}")
@@ -412,18 +417,26 @@ def process_diary(day, date, total_people, total_machinery, machinery_types, act
             detected_activities = analyze_video_activities(media_path)
             annotated_video_path = annotate_video_with_bboxes(media_path)
 
-        # Clean up temporary file if downloaded from Azure
-        if media_source == "Azure Blob" and os.path.exists(media_path):
-            os.remove(media_path)
-
         detected_types_str = ", ".join([f"{k}: {v}" for k, v in detected_machinery_types.items()])
         return [day, date, str(detected_people), str(detected_machinery), detected_types_str, detected_activities, annotated_video_path]
+
     except Exception as e:
         print(f"Error processing media: {str(e)}")
-        if media_source == "Azure Blob" and os.path.exists(media_path):
-            os.remove(media_path)
         return [day, date, "Error processing media", "Error processing media", "Error processing media", "Error processing media", None]
 
+    finally:
+        # Clean up GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            gc.collect()
+        # Remove temporary file if downloaded from Azure
+        if media_source == "Azure Blob" and media_path and os.path.exists(media_path):
+            try:
+                os.remove(media_path)
+                print(f"Removed temporary file: {media_path}")
+            except Exception as e:
+                print(f"Error removing temporary file: {str(e)}")
+
 with gr.Blocks(title="Digital Site Diary") as demo:
     gr.Markdown("# 📝 Digital Site Diary")
     with gr.Row():