Update app.py

app.py

@@ -9,15 +9,23 @@ import time
 import numpy as np
 import logging
 from glob import glob
+import zipfile
+import io
+import shutil
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Emotion mapping with emojis
-EMOJI_MAP = {
-    "angry": "😠", "disgust": "🤢", "fear": "😨",
-    "happy": "😄", "sad": "😢", "surprise": "😲", "neutral": "😐"
+# Emotion mapping with emojis and corresponding folder names
+EMOTION_MAP = {
+    "angry": {"emoji": "😠", "folder": "angry"},
+    "disgust": {"emoji": "🤢", "folder": "disgust"},
+    "fear": {"emoji": "😨", "folder": "fear"},
+    "happy": {"emoji": "😄", "folder": "happy"},
+    "sad": {"emoji": "😢", "folder": "sad"},
+    "surprise": {"emoji": "😲", "folder": "surprise"},
+    "neutral": {"emoji": "😐", "folder": "neutral"}
 }
 
 # Global variable for processing throttle
@@ -27,18 +35,24 @@ last_process_time = 0
 BACKENDS = ['mtcnn', 'opencv', 'ssd', 'dlib']
 
 # Directory to save faces
-SAVE_DIR = "/tmp/faces"
-os.makedirs(SAVE_DIR, exist_ok=True)
+BASE_SAVE_DIR = "/tmp/emotion_dataset"
+os.makedirs(BASE_SAVE_DIR, exist_ok=True)
 
-def clean_old_files(directory, max_files=20):
+# Create emotion subdirectories
+for emotion in EMOTION_MAP.values():
+    os.makedirs(os.path.join(BASE_SAVE_DIR, emotion["folder"]), exist_ok=True)
+
+def clean_old_files(directory, max_files_per_emotion=50):
     """Keep only the most recent files to avoid storage issues"""
-    files = glob(os.path.join(directory, "detected_face_*.jpg"))
-    files.sort(key=os.path.getmtime)
-    while len(files) > max_files:
-        try:
-            os.remove(files.pop(0))
-        except Exception as e:
-            logger.warning(f"Could not remove old file: {e}")
+    for emotion in EMOTION_MAP.values():
+        emotion_dir = os.path.join(directory, emotion["folder"])
+        files = glob(os.path.join(emotion_dir, "*.jpg"))
+        files.sort(key=os.path.getmtime)
+        while len(files) > max_files_per_emotion:
+            try:
+                os.remove(files.pop(0))
+            except Exception as e:
+                logger.warning(f"Could not remove old file: {e}")
 
 def enhance_image(frame):
     """Enhance image contrast to improve face detection in poor lighting."""
@@ -55,6 +69,24 @@ def enhance_image(frame):
         logger.warning(f"Image enhancement failed: {e}")
         return frame
 
+def save_face_image(face_region, emotion, confidence):
+    """Save face image to appropriate emotion folder with metadata in filename"""
+    try:
+        emotion_data = EMOTION_MAP.get(emotion, EMOTION_MAP["neutral"])
+        folder = emotion_data["folder"]
+        timestamp = int(time.time())
+        
+        # Create filename with emotion, confidence, and timestamp
+        filename = f"{emotion}_{confidence:.1f}%_{timestamp}.jpg"
+        save_path = os.path.join(BASE_SAVE_DIR, folder, filename)
+        
+        cv2.imwrite(save_path, cv2.cvtColor(face_region, cv2.COLOR_RGB2BGR))
+        logger.info(f"Saved {emotion} face to: {save_path}")
+        return save_path
+    except Exception as e:
+        logger.error(f"Failed to save face: {e}")
+        return None
+
 def predict_emotion(input_data):
     global last_process_time
     
@@ -128,29 +160,23 @@ def predict_emotion(input_data):
                 y2 = min(frame.shape[0], y + h + padding)
                 face_region = frame[y1:y2, x1:x2]
                 
-                # Save face to file if valid region
+                # Save face to appropriate emotion folder if valid region
                 if face_region.size > 0:
-                    timestamp = int(time.time())
-                    face_filename = f"detected_face_{timestamp}_{i}.jpg"
-                    save_path = os.path.join(SAVE_DIR, face_filename)
-                    try:
-                        cv2.imwrite(save_path, cv2.cvtColor(face_region, cv2.COLOR_RGB2BGR))
-                        logger.info(f"Saved face to: {save_path}")
-                    except Exception as e:
-                        logger.error(f"Failed to save face: {e}")
+                    save_face_image(face_region, emotion, confidence)
                 
                 # Draw bounding box
                 cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
                 
                 # Put emotion text on image
-                cv2.putText(frame, f"{emotion} {EMOJI_MAP.get(emotion, '')}", 
+                emoji = EMOTION_MAP.get(emotion, {}).get("emoji", "")
+                cv2.putText(frame, f"{emotion} {emoji}", 
                            (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
                 
                 # Add to text output
-                emotion_text += f"{emotion.title()} {EMOJI_MAP.get(emotion, '')}: {confidence:.1f}%\n"
+                emotion_text += f"{emotion.title()} {emoji}: {confidence:.1f}%\n"
         
         # Clean up old files
-        clean_old_files(SAVE_DIR)
+        clean_old_files(BASE_SAVE_DIR)
         
         if not emotion_text:
             emotion_text = "No faces detected. Ensure the face is clearly visible."
@@ -165,19 +191,57 @@ def predict_emotion(input_data):
             return None, f"Error: {str(e)}. Check image format or library installation."
 
 def get_saved_faces():
-    """Get list of saved face images"""
+    """Get list of saved face images from all emotion folders"""
     face_files = []
-    for f in os.listdir(SAVE_DIR):
-        if f.startswith("detected_face_") and f.endswith(".jpg"):
-            face_files.append(os.path.join(SAVE_DIR, f))
+    for emotion in EMOTION_MAP.values():
+        emotion_dir = os.path.join(BASE_SAVE_DIR, emotion["folder"])
+        for f in os.listdir(emotion_dir):
+            if f.endswith(".jpg"):
+                face_files.append(os.path.join(emotion_dir, f))
     return face_files
 
+def create_zip_file():
+    """Create a zip file containing all saved faces organized by emotion"""
+    face_files = get_saved_faces()
+    if not face_files:
+        return None
+    
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+        for face_file in face_files:
+            # Maintain directory structure in zip
+            relative_path = os.path.relpath(face_file, BASE_SAVE_DIR)
+            zip_file.write(face_file, relative_path)
+    
+    zip_buffer.seek(0)
+    return zip_buffer
+
+def create_emotion_zip(emotion):
+    """Create a zip file for a specific emotion"""
+    emotion_data = EMOTION_MAP.get(emotion)
+    if not emotion_data:
+        return None
+    
+    emotion_dir = os.path.join(BASE_SAVE_DIR, emotion_data["folder"])
+    face_files = [os.path.join(emotion_dir, f) for f in os.listdir(emotion_dir) if f.endswith(".jpg")]
+    
+    if not face_files:
+        return None
+    
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+        for face_file in face_files:
+            zip_file.write(face_file, os.path.basename(face_file))
+    
+    zip_buffer.seek(0)
+    return zip_buffer
+
 # Create Gradio interface
-with gr.Blocks(title="Emotion Detection with Face Saving") as demo:
+with gr.Blocks(title="Emotion Detection with Dataset Creation") as demo:
     gr.Markdown("""
     # Real-Time Emotion Detection
-    **Use webcam or upload an image to detect emotions**  
-    *Detected faces are saved temporarily and can be downloaded below*
+    **Use webcam or upload an image to detect emotions and build a training dataset**  
+    *Detected faces are automatically classified by emotion and saved in organized folders*
     """)
     
     with gr.Tabs():
@@ -193,46 +257,103 @@ with gr.Blocks(title="Emotion Detection with Face Saving") as demo:
             upload_text = gr.Textbox(label="Emotion Analysis")
             uploader.change(predict_emotion, uploader, [upload_output, upload_text])
         
-        with gr.TabItem("💾 Saved Faces"):
-            gr.Markdown("### Recently Detected Faces")
+        with gr.TabItem("💾 Emotion Dataset"):
+            gr.Markdown("### Classified Faces by Emotion")
+            
             with gr.Row():
-                face_gallery = gr.Gallery(label="Saved Faces", columns=4)
-                download_files = gr.File(label="Download Files", file_count="multiple")
+                with gr.Column():
+                    emotion_selector = gr.Dropdown(
+                        choices=list(EMOTION_MAP.keys()),
+                        label="Select Emotion to View/Download",
+                        value="happy"
+                    )
+                    emotion_gallery = gr.Gallery(label="Selected Emotion Faces")
+                    emotion_zip = gr.File(label="Download Emotion as ZIP", visible=False)
+                    download_emotion_btn = gr.Button("📦 Download This Emotion")
+                
+                with gr.Column():
+                    all_gallery = gr.Gallery(label="All Detected Faces")
+                    all_zip = gr.File(label="Download All as ZIP", visible=False)
+                    download_all_btn = gr.Button("📦 Download Entire Dataset")
             
             refresh_btn = gr.Button("🔄 Refresh")
             clear_btn = gr.Button("🗑️ Clear All Saved Faces")
             
-            def update_face_display():
-                faces = get_saved_faces()
-                return faces, faces
+            def update_emotion_display(emotion):
+                emotion_data = EMOTION_MAP.get(emotion, {})
+                if not emotion_data:
+                    return [], gr.File(visible=False)
+                
+                emotion_dir = os.path.join(BASE_SAVE_DIR, emotion_data["folder"])
+                face_files = [os.path.join(emotion_dir, f) for f in os.listdir(emotion_dir) if f.endswith(".jpg")]
+                return face_files, gr.File(visible=False)
+            
+            def update_all_display():
+                return get_saved_faces(), gr.File(visible=False)
+            
+            def download_selected_emotion(emotion):
+                zip_file = create_emotion_zip(emotion)
+                if zip_file:
+                    return gr.File(value=zip_file, visible=True, 
+                                 label=f"Download {emotion.capitalize()} Faces")
+                return gr.File(visible=False)
+            
+            def download_full_dataset():
+                zip_file = create_zip_file()
+                if zip_file:
+                    return gr.File(value=zip_file, visible=True, 
+                                 label="Download Entire Dataset")
+                return gr.File(visible=False)
+            
+            emotion_selector.change(
+                update_emotion_display,
+                inputs=[emotion_selector],
+                outputs=[emotion_gallery, emotion_zip]
+            )
+            
+            download_emotion_btn.click(
+                download_selected_emotion,
+                inputs=[emotion_selector],
+                outputs=[emotion_zip]
+            )
+            
+            download_all_btn.click(
+                download_full_dataset,
+                outputs=[all_zip]
+            )
             
             refresh_btn.click(
-                update_face_display,
-                outputs=[face_gallery, download_files]
+                lambda: [update_emotion_display(emotion_selector.value), 
+                        update_all_display()],
+                outputs=[emotion_gallery, emotion_zip, all_gallery, all_zip]
             )
             
             def clear_faces():
-                for f in get_saved_faces():
-                    try:
-                        os.remove(f)
-                    except:
-                        pass
-                return [], []
+                for emotion in EMOTION_MAP.values():
+                    emotion_dir = os.path.join(BASE_SAVE_DIR, emotion["folder"])
+                    for f in os.listdir(emotion_dir):
+                        if f.endswith(".jpg"):
+                            try:
+                                os.remove(os.path.join(emotion_dir, f))
+                            except:
+                                pass
+                return [], gr.File(visible=False), [], gr.File(visible=False)
             
             clear_btn.click(
                 clear_faces,
-                outputs=[face_gallery, download_files]
+                outputs=[emotion_gallery, emotion_zip, all_gallery, all_zip]
             )
     
     gr.Button("Clear All").click(
-        lambda: [None, None, "", None, None, "", [], []],
-        outputs=[webcam, webcam_output, webcam_text, uploader, upload_output, upload_text, face_gallery, download_files]
+        lambda: [None, None, "", None, None, "", [], gr.File(visible=False), [], gr.File(visible=False)],
+        outputs=[webcam, webcam_output, webcam_text, uploader, upload_output, upload_text, 
+                emotion_gallery, emotion_zip, all_gallery, all_zip]
     )
     
-    # Initialize the face display
+    # Initialize the display
     demo.load(
-        update_face_display,
-        outputs=[face_gallery, download_files]
+        lambda: [update_emotion_display(emotion_selector.value), update_all_display()],
+        outputs=[emotion_gallery, emotion_zip, all_gallery, all_zip]
     )
 
 if __name__ == "__main__":