Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import requests | |
| from transformers import pipeline | |
| # Enter your TMDb API key here | |
| TMDB_API_KEY = '2cc77b08a475b510f54a0dcca2c5c0c7' | |
| IMAGE_BASE_URL = "https://image.tmdb.org/t/p/w200" # TMDb base URL for poster images (width 200) | |
| def load_emotion_analyzer(): | |
| return pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base") | |
| emotion_analyzer = load_emotion_analyzer() | |
| def analyze_emotion(description): | |
| """Analyze emotion of a movie's description.""" | |
| emotions = emotion_analyzer(description) | |
| return emotions[0]['label'] # Return the top emotion | |
| def get_movie_genres(): | |
| """Fetches movie genres from TMDb API and returns a dictionary of genres.""" | |
| url = f'https://api.themoviedb.org/3/genre/movie/list?api_key={TMDB_API_KEY}&language=en-US' | |
| response = requests.get(url) | |
| if response.status_code == 200: | |
| genres = response.json().get('genres', []) | |
| return {genre['name']: genre['id'] for genre in genres} | |
| else: | |
| return None # Return None if the request failed | |
| def get_top_movies_by_genre(genre_id): | |
| """Fetches the top 10 highest-rated movies for a given genre ID from TMDb, sorted by title length.""" | |
| url = f'https://api.themoviedb.org/3/discover/movie?api_key={TMDB_API_KEY}&language=en-US&sort_by=vote_average.desc&vote_count.gte=1000&with_genres={genre_id}' | |
| response = requests.get(url) | |
| if response.status_code == 200: | |
| movies = response.json().get('results', [])[:10] # Get the top 10 results | |
| # Sort movies by title length, shortest to longest | |
| return sorted([(movie['title'], movie['vote_average'], movie['id'], movie['poster_path']) for movie in movies], key=lambda x: len(x[0])) | |
| else: | |
| return None | |
| def get_similar_movies(movie_id, genre_id): | |
| """Fetches similar movies for a given movie ID from TMDb and filters by the selected genre.""" | |
| url = f'https://api.themoviedb.org/3/movie/{movie_id}/similar?api_key={TMDB_API_KEY}&language=en-US' | |
| response = requests.get(url) | |
| if response.status_code == 200: | |
| movies = response.json().get('results', [])[:15] | |
| # Include the movie description (overview) in the returned values | |
| return [ | |
| (movie['title'], movie['vote_average'], movie['poster_path'], movie['id'], movie['overview']) | |
| for movie in movies if genre_id in movie['genre_ids'] and 'overview' in movie | |
| ] | |
| else: | |
| return None | |
| def get_movie_details(movie_id): | |
| """Fetches detailed information for a given movie ID from TMDb, including description, release date, director, and cast.""" | |
| url = f'https://api.themoviedb.org/3/movie/{movie_id}?api_key={TMDB_API_KEY}&language=en-US&append_to_response=credits' | |
| response = requests.get(url) | |
| if response.status_code == 200: | |
| movie = response.json() | |
| directors = [member['name'] for member in movie['credits']['crew'] if member['job'] == 'Director'] | |
| cast = [member['name'] for member in movie['credits']['cast'][:5]] # Get top 5 cast members | |
| return { | |
| "title": movie['title'], | |
| "rating": movie['vote_average'], | |
| "description": movie['overview'], | |
| "poster_path": movie['poster_path'], | |
| "release_date": movie['release_date'], | |
| "director": ", ".join(directors), | |
| "cast": ", ".join(cast) | |
| } | |
| else: | |
| return None | |
| # Initialize session state to store selected movies and recommendations | |
| if 'selected_movies' not in st.session_state: | |
| st.session_state['selected_movies'] = set() | |
| if 'recommendations' not in st.session_state: | |
| st.session_state['recommendations'] = [] | |
| if 'last_genre' not in st.session_state: | |
| st.session_state['last_genre'] = None | |
| # Streamlit app title | |
| st.title("Movie Recommender") | |
| # CSS for alignment and button styling | |
| st.markdown(""" | |
| <style> | |
| .movie-container { | |
| display: flex; | |
| flex-direction: column; | |
| align-items: center; | |
| justify-content: flex-start; | |
| } | |
| .movie-caption { | |
| min-height: 50px; /* Set a consistent height for caption sections */ | |
| display: flex; | |
| align-items: center; | |
| justify-content: center; | |
| text-align: center; | |
| } | |
| .movie-caption button { | |
| width: 100%; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # Fetch genres from TMDb and display in a dropdown | |
| genres = get_movie_genres() | |
| if genres is None: | |
| st.error("Failed to load movie genres. Please check your API key or connection.") | |
| else: | |
| genre_name = st.selectbox("Select a movie genre:", [""] + list(genres.keys())) # Include an empty option | |
| # Check if genre has changed, and reset selections if so | |
| if genre_name != st.session_state['last_genre']: | |
| st.session_state['selected_movies'].clear() | |
| st.session_state['recommendations'].clear() | |
| st.session_state['last_genre'] = genre_name | |
| # Only fetch and display movies if a genre is selected | |
| if genre_name: | |
| st.write("Please select at least one movie to generate recommendations.") | |
| genre_id = genres[genre_name] # Get the genre ID based on the selected genre name | |
| top_movies = get_top_movies_by_genre(genre_id) # Fetch top 10 movies in this genre, sorted by title length | |
| # Check if movies were successfully loaded | |
| if top_movies is None: | |
| st.error("Failed to load movies. Please check your API key or connection.") | |
| else: | |
| # Display the top 10 movies as clickable titles in a 5x2 grid layout with aligned rows | |
| st.write(f"Top 10 movies in the {genre_name} genre:") | |
| columns = st.columns(5) # Create 5 columns for each row | |
| # First row of 5 movies | |
| for index, (title, _, movie_id, poster_path) in enumerate(top_movies[:5]): | |
| col = columns[index] | |
| with col: | |
| st.markdown("<div class='movie-container'>", unsafe_allow_html=True) | |
| # Display image | |
| st.image(f"{IMAGE_BASE_URL}{poster_path}", use_container_width=True) | |
| # Display clickable title as a button to toggle selection, with consistent height | |
| if st.button(f"{title} {'✅' if movie_id in st.session_state['selected_movies'] else ''}", key=f"title_{movie_id}"): | |
| if movie_id in st.session_state['selected_movies']: | |
| st.session_state['selected_movies'].remove(movie_id) | |
| else: | |
| st.session_state['selected_movies'].add(movie_id) | |
| st.markdown("</div>", unsafe_allow_html=True) | |
| # Second row of 5 movies | |
| for index, (title, _, movie_id, poster_path) in enumerate(top_movies[5:]): | |
| col = columns[index] | |
| with col: | |
| st.markdown("<div class='movie-container'>", unsafe_allow_html=True) | |
| # Display image | |
| st.image(f"{IMAGE_BASE_URL}{poster_path}", use_container_width=True) | |
| # Display clickable title as a button to toggle selection, with consistent height | |
| if st.button(f"{title} {'✅' if movie_id in st.session_state['selected_movies'] else ''}", key=f"title_2_{movie_id}"): | |
| if movie_id in st.session_state['selected_movies']: | |
| st.session_state['selected_movies'].remove(movie_id) | |
| else: | |
| st.session_state['selected_movies'].add(movie_id) | |
| st.markdown("</div>", unsafe_allow_html=True) | |
| # Button to generate recommendations based on selected movies | |
| if st.button("Show Recommendations"): | |
| if st.session_state['selected_movies']: | |
| # Step 1: Collect emotions of user-selected movies | |
| with st.spinner("Analyzing emotions of selected movies..."): | |
| selected_emotions = [] | |
| for movie_id in st.session_state['selected_movies']: | |
| movie_details = get_movie_details(movie_id) | |
| if movie_details: | |
| emotion = analyze_emotion(movie_details["description"]) | |
| selected_emotions.append(emotion) | |
| # Step 2: Collect similar movies for all selected movies | |
| all_similar_movies = [] | |
| for movie_id in st.session_state['selected_movies']: | |
| similar_movies = get_similar_movies(movie_id, genre_id) | |
| if similar_movies: | |
| all_similar_movies.extend(similar_movies) | |
| # Step 3: Filter similar movies by matching emotions and remove duplicates | |
| recommendations = [] | |
| seen_movie_ids = set(st.session_state['selected_movies']) # Start with user-selected movies to exclude them | |
| with st.spinner("Filtering movies by description emotion..."): | |
| for title, rating, poster_path, movie_id, description in all_similar_movies: | |
| if movie_id not in seen_movie_ids: # Exclude already-seen movies | |
| movie_emotion = analyze_emotion(description) | |
| if movie_emotion in selected_emotions: | |
| seen_movie_ids.add(movie_id) # Mark this movie as seen to prevent duplicates | |
| recommendations.append((title, rating, poster_path, movie_id, movie_emotion)) | |
| # Store unique recommendations in session state | |
| st.session_state['recommendations'] = recommendations | |
| # Step 4: Display final recommendations | |
| if st.session_state['recommendations']: | |
| st.write("Based on your selections, you might also enjoy these top movies:") | |
| for title, rating, poster_path, movie_id, emotion in sorted(st.session_state['recommendations'], key=lambda x: x[1], reverse=True)[:5]: | |
| movie_details = get_movie_details(movie_id) | |
| if movie_details: | |
| col1, col2 = st.columns([1, 3]) | |
| with col1: | |
| st.image(f"{IMAGE_BASE_URL}{movie_details['poster_path']}", use_container_width=True) | |
| with col2: | |
| st.markdown(f"**Title:** {movie_details['title']}") | |
| st.markdown(f"**TMDb Rating:** {movie_details['rating']}") | |
| st.markdown(f"**Description:** {movie_details['description']}") | |
| # Additional Information button with overlay using expander | |
| with st.expander("Additional Information"): | |
| st.markdown(f"**Release Date:** {movie_details['release_date']}") | |
| st.markdown(f"**Director:** {movie_details['director']}") | |
| st.markdown(f"**Cast:** {movie_details['cast']}") | |
| else: | |
| st.warning("No movies matched your selected emotions.") |