Update app.py

app.py

@@ -1,46 +1,82 @@
+import pandas as pd
 import streamlit as st
+from surprise import Dataset, Reader, SVD
+from surprise.model_selection import train_test_split
+from collections import defaultdict
 import kagglehub
-import os
-import random
-from PIL import Image
 
 # Step 1: Download the latest version of the dataset and get the path
-path = kagglehub.dataset_download("imreallyjohn/cartoonset10k")
-images_folder = "/home/user/.cache/kagglehub/datasets/imreallyjohn/cartoonset10k/versions/1"  # Adjusted root directory
-
-# Display the dataset path
-st.write("Path to dataset files:", images_folder)
-
-# Step 2: Recursively find all image files in the dataset directory
-image_files = []
-for root, _, files in os.walk(images_folder):
-    for file in files:
-        if file.endswith('.png'):
-            image_files.append(os.path.join(root, file))  # Full path to each image
-
-# Get the count of images found
-num_images = len(image_files)
-st.write(f"Total images loaded: {num_images}")
-
-# Display an error if no images are found
-if num_images == 0:
-    st.error("No available images to display. Please check the folder path and contents.")
-else:
-    # Button to generate a random avatar if images are available
-    if st.button("Generate Random Avatar"):
-        # Step 3: Select a random available image
-        random_image_path = random.choice(image_files)  # Randomly select an image path
-        
-        # Load and display the selected image
-        selected_image = Image.open(random_image_path)
-        st.image(selected_image, caption=f"Displayed image: {os.path.basename(random_image_path)}")
-        st.write(f"Displayed image: {os.path.basename(random_image_path)}")
-        
-        # Step 4: Add a download button for the selected image
-        with open(random_image_path, "rb") as file:
-            st.download_button(
-                label="Download Avatar",
-                data=file,
-                file_name=os.path.basename(random_image_path),
-                mime="image/png"
-            )
+path = kagglehub.download_dataset("ashpalsingh1525/imdb-movies-dataset")
+
+# Step 2: Define the dataset folder path
+dataset_folder = "/home/user/.cache/kagglehub/datasets/ashpalsingh1525/imdb-movies-dataset/versions/1"
+
+# Step 3: Define the CSV file path (Update if the filename is different)
+dataset_path = f"{dataset_folder}/imdb_movies.csv"
+
+# Load the dataset
+df = pd.read_csv(dataset_path)
+
+# Ensure all categorical columns are strings
+categorical_columns = ['genre', 'orig_title', 'orig_lang', 'country', 'crew']
+
+for col in categorical_columns:
+    df[col] = df[col].astype(str)  # Convert to string explicitly
+
+# Check unique values in genre column (to ensure it's not numerical)
+if df['genre'].str.isnumeric().all():
+    print("Warning: Genre column is numeric. Mapping needed.")
+    genre_mapping = {i: f"Genre_{i}" for i in df['genre'].unique()}
+    df['genre'] = df['genre'].map(genre_mapping)
+
+# Prepare dataset for Surprise
+reader = Reader(rating_scale=(df['score'].min(), df['score'].max()))
+data = Dataset.load_from_df(df[['orig_title', 'orig_lang', 'score']], reader)
+
+# Train collaborative filtering model
+trainset, testset = train_test_split(data, test_size=0.2, random_state=42)
+model = SVD(n_factors=50, random_state=42)
+model.fit(trainset)
+
+# Function to get movie recommendations
+def get_recommendations(selected_movies, genre):
+    if not selected_movies:
+        return ["Please select at least one movie."]
+
+    # Filter dataset by genre
+    filtered_movies = df[df['genre'] == genre]
+
+    # Store average scores of all movies
+    movie_scores = defaultdict(float)
+
+    # Predict ratings for all movies in the filtered dataset
+    for movie in filtered_movies['orig_title'].unique():
+        est_score = model.predict(uid='user', iid=movie).est
+        movie_scores[movie] = est_score
+
+    # Sort movies by predicted score (descending)
+    recommended_movies = sorted(movie_scores.items(), key=lambda x: x[1], reverse=True)
+
+    # Exclude already selected movies
+    recommended_movies = [movie for movie, _ in recommended_movies if movie not in selected_movies]
+
+    return recommended_movies[:5]  # Return top 5 recommendations
+
+# Streamlit UI
+st.title("🎬 Movie Recommendation System")
+
+# Genre selection
+selected_genre = st.selectbox("Select a Genre", sorted(df['genre'].unique().tolist()))
+
+# Get available movies for the selected genre
+movies_in_genre = df[df['genre'] == selected_genre]['orig_title'].unique().tolist()
+
+# Movie selection
+selected_movies = st.multiselect("Select Up to 3 Movies", movies_in_genre, max_selections=3)
+
+# Recommendation button
+if st.button("Get Recommendations"):
+    recommendations = get_recommendations(selected_movies, selected_genre)
+    st.subheader("Recommended Movies:")
+    for movie in recommendations:
+        st.write(f"- {movie}")