Update app.py

app.py

@@ -26,22 +26,25 @@ BACKENDS = ['opencv', 'mtcnn', 'ssd', 'dlib']
 SAVE_DIR = Path("/tmp/emotion_results")
 SAVE_DIR.mkdir(exist_ok=True)
 
-# Create emotion subdirectories
+# Create directories
+(SAVE_DIR / "faces").mkdir(exist_ok=True)  # For raw face crops
+(SAVE_DIR / "annotated").mkdir(exist_ok=True)  # For annotated images
 for emotion in EMOTION_MAP.keys():
-    (SAVE_DIR / emotion).mkdir(exist_ok=True)
+    (SAVE_DIR / "faces" / emotion).mkdir(exist_ok=True, parents=True)
+    (SAVE_DIR / "annotated" / emotion).mkdir(exist_ok=True, parents=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", "image_path"])
+        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "face_path", "annotated_path"])
 
-def log_emotion(batch_no, emotion, confidence, image_path):
+def log_emotion(batch_no, emotion, confidence, face_path, annotated_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, image_path])
+        writer.writerow([timestamp, batch_no, emotion, confidence, str(face_path), str(annotated_path)])
 
 def predict_emotion(batch_no: str, image):
     if not batch_no.strip():
@@ -80,23 +83,31 @@ def predict_emotion(batch_no: str, image):
         confidence = result['emotion'][emotion]
         region = result['region']
         
-        # Draw annotations on the image
+        # Extract face coordinates
         x, y, w, h = region['x'], region['y'], region['w'], region['h']
-        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)
         
-        # Save the result in emotion-specific folder
+        # 1. Save raw face crop (for training)
+        face_crop = frame[y:y+h, x:x+w]
         timestamp = int(time.time())
-        emotion_dir = SAVE_DIR / emotion
-        output_path = str(emotion_dir / f"{batch_no}_{timestamp}.jpg")
-        cv2.imwrite(output_path, frame)
+        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}%", 
+                   (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
         
-        # Log the result
-        log_emotion(batch_no, emotion, confidence, output_path)
+        annotated_dir = SAVE_DIR / "annotated" / emotion
+        annotated_path = annotated_dir / f"{batch_no}_{timestamp}.jpg"
+        cv2.imwrite(str(annotated_path), annotated_frame)
+        
+        # Log both paths
+        log_emotion(batch_no, emotion, confidence, face_path, annotated_path)
         
         # Convert back to PIL format for display
-        output_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        output_img = Image.fromarray(cv2.cvtColor(annotated_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:
@@ -107,21 +118,66 @@ 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 create_zip(emotion=None):
-    zip_path = SAVE_DIR / "emotion_results.zip"
+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"
     with zipfile.ZipFile(zip_path, 'w') as zipf:
-        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)
+        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=[]
+    )
 
 def get_logs():
     if LOG_FILE.exists():
@@ -139,12 +195,15 @@ def download_logs():
         return str(LOG_FILE)
     return None
 
-def clear_data():
+def clear_all_data():
+    """Clear all images and logs"""
     # Clear all images
     for emotion in EMOTION_MAP.keys():
-        folder = SAVE_DIR / emotion
-        for file in folder.glob("*"):
-            file.unlink()
+        for img_type in ["faces", "annotated"]:
+            folder = SAVE_DIR / img_type / emotion
+            for file in folder.glob("*"):
+                if file.is_file():
+                    file.unlink()
     
     # Clear logs
     if LOG_FILE.exists():
@@ -153,7 +212,7 @@ def clear_data():
     # Recreate empty log file
     with open(LOG_FILE, 'w', newline='') as f:
         writer = csv.writer(f)
-        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "image_path"])
+        writer.writerow(["timestamp", "batch_no", "emotion", "confidence", "face_path", "annotated_path"])
     
     return "All data has been cleared", pd.DataFrame(), None
 
@@ -161,6 +220,7 @@ def clear_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("""
@@ -205,44 +265,117 @@ with gr.Blocks(title="Emotion Capture", css="""
 
 # Data Management Interface
 with gr.Blocks(title="Data Management", css="""
-    .gradio-container { max-width: 900px !important }
+    .gradio-container { max-width: 1200px !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("Download Images"):
-        with gr.Column(elem_classes="data-section"):
-            gr.Markdown("## Download Options")
+    with gr.Tab("Image Management"):
+        with gr.Column():
+            gr.Markdown("## Select and Manage Images")
             with gr.Row():
                 emotion_selector = gr.Dropdown(
                     choices=["All Emotions"] + list(EMOTION_MAP.keys()),
-                    label="Select Emotion Category",
-                    value="All Emotions",
-                    interactive=True
+                    label="Emotion Category",
+                    value="All Emotions"
+                )
+                image_type_selector = gr.Dropdown(
+                    choices=["faces", "annotated"],
+                    label="Image Type",
+                    value="faces"
                 )
-                download_btn = gr.Button("Download Images as ZIP")
-            download_file = gr.File(label="Download Result")
+                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")
             
+            # 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_zip,
-                inputs=emotion_selector,
-                outputs=download_file
+                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]
             )
     
     with gr.Tab("Emotion Logs"):
-        with gr.Column(elem_classes="data-section"):
-            gr.Markdown("## Emotion Logs")
+        with gr.Column():
+            gr.Markdown("## Emotion Analysis Logs")
             with gr.Row():
-                refresh_btn = gr.Button("Refresh Logs")
+                refresh_logs_btn = gr.Button("Refresh Logs")
                 download_logs_btn = gr.Button("Download Logs as CSV")
-                clear_btn = gr.Button("Clear All Data", variant="stop")
+                clear_all_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_btn.click(
+            refresh_logs_btn.click(
                 view_logs,
                 outputs=logs_display
             )
@@ -252,9 +385,12 @@ with gr.Blocks(title="Data Management", css="""
                 outputs=logs_csv
             )
             
-            clear_btn.click(
-                clear_data,
+            clear_all_btn.click(
+                clear_all_data,
                 outputs=[clear_message, logs_display, logs_csv]
+            ).then(
+                lambda: update_gallery("All Emotions", "faces"),
+                outputs=[gallery, selected_images]
             )
             
             # Initial load of logs
@@ -266,7 +402,16 @@ with gr.Blocks(title="Data Management", css="""
 # Combine interfaces
 demo = gr.TabbedInterface(
     [capture_interface, data_interface],
-    ["Emotion Capture", "Data Management"]
+    ["Emotion Capture", "Data Management"],
+    css="""
+    .tab { padding: 20px; }
+    """
+)
+
+# Initialize gallery and logs
+demo.load(
+    lambda: update_gallery("All Emotions", "faces"),
+    outputs=[gallery, selected_images]
 )
 
 if __name__ == "__main__":