Update app.py

app.py

@@ -26,25 +26,22 @@ BACKENDS = ['opencv', 'mtcnn', 'ssd', 'dlib']
 SAVE_DIR = Path("/tmp/emotion_results")
 SAVE_DIR.mkdir(exist_ok=True)
 
-# Create directories
-(SAVE_DIR / "faces").mkdir(exist_ok=True)  # For raw face crops
-(SAVE_DIR / "annotated").mkdir(exist_ok=True)  # For annotated images
+# Create emotion subdirectories
 for emotion in EMOTION_MAP.keys():
-    (SAVE_DIR / "faces" / emotion).mkdir(exist_ok=True, parents=True)
-    (SAVE_DIR / "annotated" / emotion).mkdir(exist_ok=True, parents=True)
+    (SAVE_DIR / emotion).mkdir(exist_ok=True)
 
 # Log file setup
 LOG_FILE = SAVE_DIR / "emotion_logs.csv"
 if not LOG_FILE.exists():
     with open(LOG_FILE, 'w', newline='') as f:
         writer = csv.writer(f)
-        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "face_path", "annotated_path"])
+        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "image_path"])
 
-def log_emotion(batch_no, emotion, confidence, face_path, annotated_path):
+def log_emotion(batch_no, emotion, confidence, image_path):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     with open(LOG_FILE, 'a', newline='') as f:
         writer = csv.writer(f)
-        writer.writerow([timestamp, batch_no, emotion, confidence, str(face_path), str(annotated_path)])
+        writer.writerow([timestamp, batch_no, emotion, confidence, image_path])
 
 def predict_emotion(batch_no: str, image):
     if not batch_no.strip():
@@ -83,31 +80,23 @@ def predict_emotion(batch_no: str, image):
         confidence = result['emotion'][emotion]
         region = result['region']
         
-        # Extract face coordinates
+        # Draw annotations on the image
         x, y, w, h = region['x'], region['y'], region['w'], region['h']
-        
-        # 1. Save raw face crop (for training)
-        face_crop = frame[y:y+h, x:x+w]
-        timestamp = int(time.time())
-        face_dir = SAVE_DIR / "faces" / emotion
-        face_path = face_dir / f"{batch_no}_{timestamp}.jpg"
-        cv2.imwrite(str(face_path), face_crop)
-        
-        # 2. Create and save annotated image (for display)
-        annotated_frame = frame.copy()
-        cv2.rectangle(annotated_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
-        cv2.putText(annotated_frame, f"{emotion} {EMOTION_MAP[emotion]} {confidence:.1f}%", 
+        cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
+        cv2.putText(frame, f"{emotion} {EMOTION_MAP[emotion]} {confidence:.1f}%", 
                    (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
         
-        annotated_dir = SAVE_DIR / "annotated" / emotion
-        annotated_path = annotated_dir / f"{batch_no}_{timestamp}.jpg"
-        cv2.imwrite(str(annotated_path), annotated_frame)
+        # Save the result in emotion-specific folder
+        timestamp = int(time.time())
+        emotion_dir = SAVE_DIR / emotion
+        output_path = str(emotion_dir / f"{batch_no}_{timestamp}.jpg")
+        cv2.imwrite(output_path, frame)
         
-        # Log both paths
-        log_emotion(batch_no, emotion, confidence, face_path, annotated_path)
+        # Log the result
+        log_emotion(batch_no, emotion, confidence, output_path)
         
         # Convert back to PIL format for display
-        output_img = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
+        output_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
         return output_img, f"Batch {batch_no}: {emotion.title()} ({confidence:.1f}%)", "", gr.Image(visible=True), gr.Textbox(visible=True)
     
     except Exception as e:
@@ -118,66 +107,21 @@ def toggle_webcam(batch_no):
         return gr.Image(visible=True), gr.Image(visible=False), gr.Textbox(visible=False), ""
     return gr.Image(visible=False), gr.Image(visible=False), gr.Textbox(visible=False), "Please enter a batch number"
 
-def get_image_gallery(emotion, image_type):
-    """Get image gallery for selected emotion and type"""
-    if emotion == "All Emotions":
-        image_dict = {}
-        for emot in EMOTION_MAP.keys():
-            folder = SAVE_DIR / image_type / emot
-            image_dict[emot] = [str(f) for f in folder.glob("*.jpg") if f.exists()]
-    else:
-        folder = SAVE_DIR / image_type / emotion
-        image_dict = {emotion: [str(f) for f in folder.glob("*.jpg") if f.exists()]}
-    return image_dict
-
-def create_custom_zip(selected_images):
-    """Create zip from selected images"""
-    if not selected_images:
-        return None
-    zip_path = SAVE_DIR / "selected_images.zip"
+def create_zip(emotion=None):
+    zip_path = SAVE_DIR / "emotion_results.zip"
     with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for img_path in selected_images:
-            if Path(img_path).exists():
-                zipf.write(img_path, arcname=Path(img_path).name)
-    return str(zip_path) if zip_path.exists() else None
-
-def delete_selected_images(selected_images):
-    """Delete selected images"""
-    if not selected_images:
-        return "No images selected"
-    
-    deleted_count = 0
-    for img_path in selected_images:
-        try:
-            Path(img_path).unlink()
-            deleted_count += 1
-        except Exception as e:
-            print(f"Error deleting {img_path}: {e}")
-    
-    # Update logs
-    if LOG_FILE.exists():
-        df = pd.read_csv(LOG_FILE)
-        for img_path in selected_images:
-            if "faces" in img_path:
-                df = df[df.face_path != img_path]
-            else:
-                df = df[df.annotated_path != img_path]
-        df.to_csv(LOG_FILE, index=False)
-    
-    return f"Deleted {deleted_count} images"
-
-def update_gallery(emotion, image_type):
-    """Update the image gallery view"""
-    image_dict = get_image_gallery(emotion, image_type)
-    gallery = []
-    for emotion, images in image_dict.items():
-        for img_path in images:
-            gallery.append((img_path, f"{emotion}: {Path(img_path).name}"))
-    return gr.Gallery(value=gallery, label="Image Gallery"), gr.CheckboxGroup(
-        choices=[img[0] for img in gallery], 
-        label="Selected Images",
-        value=[]
-    )
+        if emotion and emotion != "All Emotions":
+            # Zip specific emotion folder
+            folder = SAVE_DIR / emotion
+            for file in folder.iterdir():
+                zipf.write(file, arcname=f"{emotion}/{file.name}")
+        else:
+            # Zip all images
+            for emotion in EMOTION_MAP.keys():
+                folder = SAVE_DIR / emotion
+                for file in folder.iterdir():
+                    zipf.write(file, arcname=f"{emotion}/{file.name}")
+    return str(zip_path)
 
 def get_logs():
     if LOG_FILE.exists():
@@ -195,15 +139,12 @@ def download_logs():
         return str(LOG_FILE)
     return None
 
-def clear_all_data():
-    """Clear all images and logs"""
+def clear_data():
     # Clear all images
     for emotion in EMOTION_MAP.keys():
-        for img_type in ["faces", "annotated"]:
-            folder = SAVE_DIR / img_type / emotion
-            for file in folder.glob("*"):
-                if file.is_file():
-                    file.unlink()
+        folder = SAVE_DIR / emotion
+        for file in folder.glob("*"):
+            file.unlink()
     
     # Clear logs
     if LOG_FILE.exists():
@@ -212,7 +153,7 @@ def clear_all_data():
     # Recreate empty log file
     with open(LOG_FILE, 'w', newline='') as f:
         writer = csv.writer(f)
-        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "face_path", "annotated_path"])
+        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "image_path"])
     
     return "All data has been cleared", pd.DataFrame(), None
 
@@ -220,7 +161,6 @@ def clear_all_data():
 with gr.Blocks(title="Emotion Capture", css="""
     .gradio-container { max-width: 800px !important }
     .message { color: red; font-weight: bold; }
-    .gallery { grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); }
 """) as capture_interface:
     
     gr.Markdown("""
@@ -265,117 +205,44 @@ with gr.Blocks(title="Emotion Capture", css="""
 
 # Data Management Interface
 with gr.Blocks(title="Data Management", css="""
-    .gradio-container { max-width: 1200px !important }
+    .gradio-container { max-width: 900px !important }
     .data-section { border: 1px solid #ccc; padding: 20px; border-radius: 5px; margin-top: 20px; }
-    .gallery { grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); }
-    .action-buttons { margin-top: 20px; }
 """) as data_interface:
     
     gr.Markdown("# Data Management Interface")
     
-    with gr.Tab("Image Management"):
-        with gr.Column():
-            gr.Markdown("## Select and Manage Images")
+    with gr.Tab("Download Images"):
+        with gr.Column(elem_classes="data-section"):
+            gr.Markdown("## Download Options")
             with gr.Row():
                 emotion_selector = gr.Dropdown(
                     choices=["All Emotions"] + list(EMOTION_MAP.keys()),
-                    label="Emotion Category",
-                    value="All Emotions"
-                )
-                image_type_selector = gr.Dropdown(
-                    choices=["faces", "annotated"],
-                    label="Image Type",
-                    value="faces"
+                    label="Select Emotion Category",
+                    value="All Emotions",
+                    interactive=True
                 )
-                refresh_btn = gr.Button("Refresh Gallery")
-            
-            gallery = gr.Gallery(
-                label="Image Gallery",
-                elem_classes="gallery",
-                columns=4
-            )
-            selected_images = gr.CheckboxGroup(
-                label="Selected Images",
-                interactive=True
-            )
-            
-            with gr.Row(variant="panel"):
-                with gr.Column():
-                    gr.Markdown("### Download Options")
-                    download_btn = gr.Button("Download Selected", variant="primary")
-                    download_all_btn = gr.Button("Download All in Category")
-                    download_output = gr.File(label="Download Result")
-                
-                with gr.Column():
-                    gr.Markdown("### Delete Options")
-                    delete_btn = gr.Button("Delete Selected", variant="stop")
-                    delete_all_btn = gr.Button("Delete All in Category", variant="stop")
-                    delete_output = gr.Textbox(label="Delete Status")
+                download_btn = gr.Button("Download Images as ZIP")
+            download_file = gr.File(label="Download Result")
             
-            # Update gallery when parameters change
-            emotion_selector.change(
-                update_gallery,
-                inputs=[emotion_selector, image_type_selector],
-                outputs=[gallery, selected_images]
-            )
-            
-            image_type_selector.change(
-                update_gallery,
-                inputs=[emotion_selector, image_type_selector],
-                outputs=[gallery, selected_images]
-            )
-            
-            refresh_btn.click(
-                update_gallery,
-                inputs=[emotion_selector, image_type_selector],
-                outputs=[gallery, selected_images]
-            )
-            
-            # Download handlers
             download_btn.click(
-                create_custom_zip,
-                inputs=selected_images,
-                outputs=download_output
-            )
-            
-            download_all_btn.click(
-                lambda e, t: create_custom_zip([img[0] for img in get_image_gallery(e, t).items() for img in img[1]]),
-                inputs=[emotion_selector, image_type_selector],
-                outputs=download_output
-            )
-            
-            # Delete handlers
-            delete_btn.click(
-                delete_selected_images,
-                inputs=selected_images,
-                outputs=delete_output
-            ).then(
-                lambda: update_gallery(emotion_selector.value, image_type_selector.value),
-                outputs=[gallery, selected_images]
-            )
-            
-            delete_all_btn.click(
-                lambda e, t: delete_selected_images([img for img_list in get_image_gallery(e, t).values() for img in img_list]),
-                inputs=[emotion_selector, image_type_selector],
-                outputs=delete_output
-            ).then(
-                lambda: update_gallery(emotion_selector.value, image_type_selector.value),
-                outputs=[gallery, selected_images]
+                create_zip,
+                inputs=emotion_selector,
+                outputs=download_file
             )
     
     with gr.Tab("Emotion Logs"):
-        with gr.Column():
-            gr.Markdown("## Emotion Analysis Logs")
+        with gr.Column(elem_classes="data-section"):
+            gr.Markdown("## Emotion Logs")
             with gr.Row():
-                refresh_logs_btn = gr.Button("Refresh Logs")
+                refresh_btn = gr.Button("Refresh Logs")
                 download_logs_btn = gr.Button("Download Logs as CSV")
-                clear_all_btn = gr.Button("Clear All Data", variant="stop")
+                clear_btn = gr.Button("Clear All Data", variant="stop")
             
             logs_display = gr.Markdown()
             logs_csv = gr.File(label="Logs Download")
             clear_message = gr.Textbox(label="Status", interactive=False)
             
-            refresh_logs_btn.click(
+            refresh_btn.click(
                 view_logs,
                 outputs=logs_display
             )
@@ -385,12 +252,9 @@ with gr.Blocks(title="Data Management", css="""
                 outputs=logs_csv
             )
             
-            clear_all_btn.click(
-                clear_all_data,
+            clear_btn.click(
+                clear_data,
                 outputs=[clear_message, logs_display, logs_csv]
-            ).then(
-                lambda: update_gallery("All Emotions", "faces"),
-                outputs=[gallery, selected_images]
             )
             
             # Initial load of logs
@@ -402,16 +266,7 @@ with gr.Blocks(title="Data Management", css="""
 # Combine interfaces
 demo = gr.TabbedInterface(
     [capture_interface, data_interface],
-    ["Emotion Capture", "Data Management"],
-    css="""
-    .tab { padding: 20px; }
-    """
-)
-
-# Initialize gallery and logs
-demo.load(
-    lambda: update_gallery("All Emotions", "faces"),
-    outputs=[gallery, selected_images]
+    ["Emotion Capture", "Data Management"]
 )
 
 if __name__ == "__main__":