Refactor with GPT
- Dominant language
- Python
- Stars
- 6
- Forks
- 2
- PR merge metrics
- No merged PRs in 30d
Description
Sure! Let's refactor the code using object-oriented design principles, ensure proper logging, and provide thorough docstrings. Here are the classes and integration pieces as requested:
1. An `EmotionDetector` class for facial emotion detection.
2. An `EmotionStorage` class for handling interactions with the SQLite database.
3. An `EmotionVisualizer` class for visualizing emotions using graphs.
4. A main `EmoTrackApp` class that coordinates these components.
5. Logging and argparse for command-line arguments.
Here is the refactored code:
```python
import logging
import sqlite3
from datetime import datetime
import boto3
import cv2
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import streamlit as st
from argparse import ArgumentParser
# Configuration
BATCH_SIZE = 60
# Initialize logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
class EmotionDetector:
"""Class for detecting emotions from facial analysis."""
def __init__(self):
"""Initialize the EmotionDetector."""
self.client = boto3.client("rekognition")
logging.info("EmotionDetector initialized")
def detect_emotion(self, frame):
"""Detects the emotion of a face in a frame."""
ret, jpg_data = cv2.imencode(".jpg", frame)
if not ret:
logging.error("Failed to encode frame")
raise ValueError("Failed to encode frame")
image_bytes = jpg_data.tobytes()
response = self.client.detect_faces(
Image={"Bytes": image_bytes},
Attributes=["ALL"]
)
if not response['FaceDetails']:
return "NO FACE"
emotions = response['FaceDetails'][0]['Emotions']
emotions.sort(key=lambda x: x['Confidence'], reverse=True)
return emotions[0]['Type'] if emotions else "UNKNOWN"
class EmotionStorage:
"""Class for storing and retrieving emotions in a SQLite database."""
def __init__(self, db_path="emotions.db"):
"""Initialize the EmotionStorage."""
self.db_path = db_path
self._initialize_db()
logging.info("EmotionStorage initialized with database: %s", db_path)
def _initialize_db(self):
"""Initialize the SQLite database and create the emotions table if it doesn't exist."""
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute(
"""CREATE TABLE IF NOT EXISTS emotions
(timestamp INTEGER, emotion TEXT)"""
)
conn.commit()
logging.info("Database initialized and table created if not exists")
def save_emotions_batch(self, emotions_batch):
"""Save a batch of emotions to the SQLite database."""
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.executemany(
"INSERT INTO emotions (timestamp, emotion) VALUES (?, ?)", emotions_batch
)
conn.commit()
logging.info("Saved batch of emotions to the database")
class EmotionVisualizer:
"""Class for visualizing emotions using graphs."""
def __init__(self, db_path="emotions.db"):
"""Initialize the EmotionVisualizer."""
self.db_path = db_path
logging.info("EmotionVisualizer initialized with database: %s", db_path)
def display_emotion_graphs(self):
"""Display various graphs illustrating the emotions."""
with sqlite3.connect(self.db_path) as conn:
query = """
SELECT
DATE(DATETIME(timestamp, 'unixepoch')) AS date,
emotion,
COUNT(emotion) AS emotion_count
FROM
emotions
GROUP BY
date, emotion
ORDER BY
date DESC, emotion_count DESC;
"""
df = pd.read_sql_query(query, conn)
if df.empty:
st.warning("No data available to display.")
return
total_counts = df.groupby("date")["emotion_count"].sum().reset_index()
total_counts.rename(columns={"emotion_count": "total_count"}, inplace=True)
df = pd.merge(df, total_counts, on="date")
df["percentage"] = (df["emotion_count"] / df["total_count"]) * 100
self._plot_stacked_bar(df)
self._plot_line_graph(df)
def _plot_stacked_bar(self, df):
"""Plot a stacked bar graph showing emotion variation."""
fig, ax = plt.subplots(figsize=(14, 6))
bottom_values = {date: 0 for date in df["date"].unique()}
emotion_colors = {
"CALM": "#D3D3D3",
"SURPRISED": "#A9A9A9",
"CONFUSED": "#808080",
"HAPPY": "#696969",
"SAD": "#800000",
"ANGRY": "#8B0000",
"FEAR": "#A52A2A",
}
for emotion in df["emotion"].unique():
emotion_data = df[df["emotion"] == emotion]
bottoms = [bottom_values[date] for date in emotion_data["date"]]
bars = ax.bar(
emotion_data["date"],
emotion_data["percentage"],
bottom=bottoms,
color=emotion_colors.get(emotion, "white"),
)
for i, date in enumerate(emotion_data["date"]):
bottom_values[date] += emotion_data.iloc[i]["percentage"]
for bar, percentage in zip(bars, emotion_data["percentage"]):
height = bar.get_height()
position = bar.get_y()
if height > 0:
text_color = "black" if emotion in ["CALM", "SURPRISED", "CONFUSED", "HAPPY"] else "white"
ax.text(
bar.get_x() + bar.get_width() / 2,
position + height / 2,
emotion,
ha="center",
va="center",
color=text_color,
)
ax.set_title("Emotion Variation in the Past 7 Days (100% Stacked)")
ax.set_xticklabels(df["date"].unique(), rotation=45)
plt.legend(df["emotion"].unique())
st.pyplot(fig)
def _plot_line_graph(self, df):
"""Plot a line graph showing emotion trends over time."""
fig, ax = plt.subplots(figsize=(14, 6))
for emotion in df["emotion"].unique():
emotion_data = df[df["emotion"] == emotion]
ax.plot(emotion_data["date"], emotion_data["emotion_count"], marker='o', label=emotion)
ax.set_title("Emotion Trends Over Time")
ax.set_xlabel("Date")
ax.set_ylabel("Number of Emotions")
ax.legend()
st.pyplot(fig)
class EmoTrackApp:
"""Main class for the EmoTrack application."""
def __init__(self, title="EmoTrack"):
"""Initialize the EmoTrackApp with necessary components."""
self.title = title
self.detector = EmotionDetector()
self.storage = EmotionStorage()
self.visualizer = EmotionVisualizer()
st.set_page_config(page_title=title)
logging.info("EmoTrackApp initialized")
def run(self):
"""Run the EmoTrack application."""
st.title(self.title)
self._run_webcam_feed()
self._display_graphs()
def _run_webcam_feed(self):
"""Run the webcam feed for detecting and saving emotions."""
st.write("## Webcam Feed")
if "running" not in st.session_state:
st.session_state.running = False
if st.button("Start", key="start_button"):
st.session_state.running = True
if st.button("Stop", key="stop_button"):
st.session_state.running = False
frame_slot = st.empty()
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
frame_count = 0
current_emotion = None
emotions_batch = []
while st.session_state.running:
ret, frame = cap.read()
if not ret:
st.warning("Failed to get frame from webcam.")
st.session_state.running = False
break
frame_count += 1
if frame_count % 24 == 0:
current_emotion = self.detector.detect_emotion(frame)
if current_emotion != "NO FACE":
emotions_batch.append(
(
datetime.now().timestamp(),
current_emotion,
)
)
if len(emotions_batch) == BATCH_SIZE:
self.storage.save_emotions_batch(emotions_batch)
emotions_batch = []
if current_emotion:
cv2.putText(
frame,
current_emotion,
(10, 30),
cv2.FONT_HERSHEY_SIMPLEX,
0.9,
(0, 255, 0),
2,
)
frame_slot.image(frame, channels="BGR", use_column_width=True)
if emotions_batch:
self.storage.save_emotions_batch(emotions_batch)
cap.release()
def _display_graphs(self):
"""Display emotion graphs."""
st.write("## Emotion Variation per Day")
self.visualizer.display_emotion_graphs()
def main():
parser = ArgumentParser(description="EmoTrack Application")
parser.add_argument("--title", type=str, default="EmoTrack", help="The title of the application.")
args = parser.parse_args()
app = EmoTrackApp(title=args.title)
app.run()
if __name__ == "__main__":
main()
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.