Update app.py

app.py

@@ -7,95 +7,81 @@ from dotenv import load_dotenv
 from google import genai
 from google.genai import types
 
-# Import our DINO processing function
-from dino_processor import process_video_with_dino
+# Import the optimized DINO functions
+from dino_processor import load_dino_model, process_video_with_dino
 
-# --- Configuration and Client Initialization ---
+# --- Configuration and Initialization ---
 load_dotenv()
 try:
-    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+    GEMINI_CLIENT = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 except KeyError:
     raise gr.Error("FATAL: GEMINI_API_KEY not found. Please set it in your Hugging Face Space secrets.")
 
+DINO_MODEL, DINO_DEVICE = load_dino_model()
+
 # --- Main Gradio Generator Function ---
 def full_analysis(video_file_path, user_observations):
     if not video_file_path:
-        # Yield a single update and stop
-        yield "Please upload a video first.", "", "", None, None
+        yield "Please upload a video first.", "", "", None
         return
 
-    # --- Stage 1: Gemini Analysis (Fast) ---
+    # Stage 1: Gemini Analysis
     print("--- Starting Device Analysis ---")
     device_type, condition, reason = "Processing...", "Processing...", "Processing..."
-    
-    # Immediately update the UI to show that Gemini is working
-    yield device_type, condition, reason, None, None
+    yield device_type, condition, reason, None
 
     uploaded_file = None
     try:
-        uploaded_file = client.files.upload(file=video_file_path)
+        uploaded_file = GEMINI_CLIENT.files.upload(file=video_file_path)
         while uploaded_file.state.name == "PROCESSING":
             time.sleep(5)
-            uploaded_file = client.files.get(name=uploaded_file.name)
+            uploaded_file = GEMINI_CLIENT.files.get(name=uploaded_file.name)
         if uploaded_file.state.name != "ACTIVE":
             raise Exception(f"File processing failed for Gemini. State: {uploaded_file.state.name}")
         
         prompt = f"""
-        Analyze the provided video and the user's observations. Respond ONLY with a valid JSON object.
+        Analyze the provided video and user observations. Respond ONLY with a valid JSON object.
         User Observations: "{user_observations if user_observations else 'None'}"
-        Follow these steps:
-        1. Identify the primary object. Is it an electronic device or a white good?
-        2. If YES, respond with a JSON object with keys: "device_type", "condition" (choose from "unusable", "poor", "functional", "good", "excellent"), and "reason" (a brief explanation).
-        3. If NO, respond with a JSON object with only ONE key: "device_type" starting with "Invalid: " (e.g., "Invalid: Water Bottle").
+        Logic:
+        1. Is the object an electronic device or white good?
+        2. If YES: respond with JSON keys "device_type", "condition" (from "unusable", "poor", "functional", "good", "excellent"), and "reason".
+        3. If NO: respond with JSON key "device_type" starting with "Invalid: ".
         """
-        
-        model_name = "gemini-2.5-flash"
+        model_name = "gemini-1.5-flash-latest" # Using a reliable public model
         config = types.GenerateContentConfig(temperature=0.2, response_mime_type="application/json")
         contents = [uploaded_file, prompt]
-        
-        response = client.models.generate_content(model=f"models/{model_name}", contents=contents, config=config)
+        response = GEMINI_CLIENT.models.generate_content(model=f"models/{model_name}", contents=contents, config=config)
         parsed_json = json.loads(response.text)
-        
         device_type = parsed_json.get("device_type", "N/A")
         if "Invalid:" in device_type:
-            condition = "N/A"
-            reason = "Object is not a valid electronic device or white good."
+            condition, reason = "N/A", "Object is not a valid device."
         else:
-            condition = parsed_json.get("condition", "N/A")
-            reason = parsed_json.get("reason", "N/A")
-            
+            condition, reason = parsed_json.get("condition", "N/A"), parsed_json.get("reason", "N/A")
         print("--- Device Analysis Complete ---")
-        
     except Exception as e:
         print(f"ERROR during Gemini processing: {e}")
         device_type, condition, reason = f"Error in Gemini processing: {e}", "", ""
     finally:
         if uploaded_file:
-            client.files.delete(name=uploaded_file.name)
+            GEMINI_CLIENT.files.delete(name=uploaded_file.name)
 
-    # --- YIELD 1: Update UI with Gemini results ---
-    # The DINO galleries are still empty (None) at this point.
-    yield device_type, condition, reason, None, None
+    yield device_type, condition, reason, None
 
-    # --- Stage 2: DINO Analysis (Slow) ---
+    # Stage 2: DINO Analysis
     print("--- Starting DINO Analysis ---")
     try:
-        dino_results = process_video_with_dino(video_file_path)
-        overlay_images = [res[0] for res in dino_results]
-        attention_maps = [res[1] for res in dino_results]
+        # OPTIMIZATION: This now returns a simple list of overlay image paths
+        overlay_images = process_video_with_dino(video_file_path, DINO_MODEL, DINO_DEVICE)
         print("--- DINO Analysis Complete ---")
     except Exception as e:
         print(f"ERROR during DINO processing: {e}")
-        # If DINO fails, we still want to keep the Gemini results.
-        # We can update one of the text fields to show the DINO error.
         reason += f" | DINO Error: {e}"
-        overlay_images, attention_maps = None, None
+        overlay_images = None
 
-    # --- YIELD 2: Final update with all results ---
-    # This updates the galleries with the DINO images.
-    yield device_type, condition, reason, overlay_images, attention_maps
+    # OPTIMIZATION: Final yield with only 4 items
+    yield device_type, condition, reason, overlay_images
 
-# --- Gradio Interface ---
+# --- Gradio Interface (Simplified) ---
 theme = gr.themes.Soft(
     primary_hue="green", secondary_hue="emerald"
 ).set(
@@ -111,10 +97,7 @@ with gr.Blocks(theme=theme) as demo:
     with gr.Row():
         with gr.Column(scale=1):
             video_input = gr.Video(label="Upload or Record Video", sources=["upload", "webcam"], format="mp4")
-            observations_input = gr.Textbox(
-                label="Optional: User Observations", 
-                placeholder="e.g., Device overheats, screen flickers..."
-            )
+            observations_input = gr.Textbox(label="Optional: User Observations", placeholder="e.g., Device overheats...")
             submit_button = gr.Button("Run Full Analysis", variant="primary")
         
         with gr.Column(scale=2):
@@ -126,26 +109,19 @@ with gr.Blocks(theme=theme) as demo:
     
     gr.Markdown("---")
     
-    gr.Markdown("## DINO: Visual Feature Analysis")
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### Overlayed Heatmaps on Key Frames")
-            overlay_gallery = gr.Gallery(label="Overlayed Heatmaps", elem_id="gallery", columns=5)
-        with gr.Column():
-            gr.Markdown("### Self-Attention Maps")
-            attention_gallery = gr.Gallery(label="Self-Attention Maps", elem_id="gallery", columns=5)
+    # OPTIMIZATION: Simplified DINO output section
+    gr.Markdown("## Overlayed Heatmaps on Key Frames")
+    overlay_gallery = gr.Gallery(label="Overlayed Heatmaps", elem_id="gallery", columns=3)
 
-    # The .click() event now points to a generator function.
-    # Gradio will automatically handle the multiple `yield` statements.
     submit_button.click(
         fn=full_analysis,
         inputs=[video_input, observations_input],
+        # OPTIMIZATION: Updated outputs list to have only 4 items
         outputs=[
             device_type_output, 
             condition_output, 
             reason_output, 
-            overlay_gallery, 
-            attention_gallery
+            overlay_gallery
         ],
         show_progress='full'
     )