Notebook for matching scenes with StashDB UUIDs with TPDB

This commit is contained in:
Minasuki Hikimuna
2024-08-07 13:08:29 +03:00
parent ca51dd3775
commit beecf04525
4 changed files with 577 additions and 0 deletions
+27
View File
@@ -44,3 +44,30 @@ After opening this repository in Visual Studio Code, remember to set correct ker
Now you can start experimenting with Stash and Pandas by looking at hello_world.ipynb.
![Example notebook with Data Wrangler](images/hello_world_data_wrangler.png)
# Notebooks
## Hello World
A handful of samples to get you going.
## Scene Matching
Tailored to a workflow where scenes have been matched to StashDB first and you want to match to TPDB without using TPDB's data directly.
Process is as follows:
1. Set TPDB_API_KEY value in .env file to your TPDB API key
2. Run initial cells to create Stash app connection and create utility functions
3. Run studio summary to figure out which networks have most scenes missing TPDB. Open `df_scenes_grouped_by_studio` in Data Wrangler.
4. Iterate following process
1. Find studio with name
1. Get 100 scenes at a time from Stash app
1. Get TPDB matches with phash and get the best match in a single table
1. Open `df_matched_scenes` in Data Wrangler
1. Manually verify with Data Wrangler UI that scenes match as you expected
1. Save TPDB UUIDs to your Stash app with one of the last cells. There are multiple examples based on if you are comfortable saving TPDB UUIDs only to those scenes which match all identifiers or if you want something more specific.
![Studios with most missing scenes](images/scene_matching_studio_summary.png)
![Matched scenes under review](images/scene_matching_manual_review.png)
Binary file not shown.

After

Width:  |  Height:  |  Size: 147 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 116 KiB

+550
View File
@@ -0,0 +1,550 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create Stash app client\n",
"\n",
"import pandas as pd\n",
"import dotenv\n",
"import os\n",
"\n",
"from libraries.client_stashapp import get_stashapp_client\n",
"\n",
"dotenv.load_dotenv()\n",
"\n",
"stash = get_stashapp_client()\n",
"\n",
"tpdb_headers = {\n",
" \"Content-Type\": \"application/json\",\n",
" \"Authorization\": f\"Bearer {os.environ['TPDB_API_KEY']}\",\n",
"}\n",
"\n",
"# Utility functions\n",
"def save_tpdb_uuid(df):\n",
" for _, row in df.iterrows():\n",
" merged_stash_ids = row[\"stash_ids\"]\n",
" merged_stash_ids.append({\n",
" \"endpoint\": \"https://theporndb.net/graphql\",\n",
" \"stash_id\": row[\"tpdb_uuid\"]\n",
" })\n",
"\n",
" # Assume update_scene returns True if successful, False otherwise\n",
" success = stash.update_scene(\n",
" {\n",
" \"id\": row[\"id\"],\n",
" \"stash_ids\": merged_stash_ids\n",
" },\n",
" False\n",
" )\n",
"\n",
" if success:\n",
" # If update was successful, add the index to the list\n",
" print(f\"Updated scene {row['id']} with TPDB ID {row['tpdb_uuid']}\")\n",
"\n",
"def is_tpdb_endpoint_missing(stash_ids):\n",
" # Ensure stash_ids is a list\n",
" if isinstance(stash_ids, list):\n",
" # Check if no dictionary in the list has the required endpoint\n",
" return not any(item.get(\"endpoint\") == \"https://theporndb.net/graphql\" for item in stash_ids)\n",
" return True # If stash_ids is not a list, assume endpoint is missing\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Find which parent studios and studios have most scenes missing TPDB IDs\n",
"scenes_with_stashdb = stash.find_scenes(\n",
" {\n",
" \"stash_id_endpoint\": {\n",
" \"endpoint\": \"https://stashdb.org/graphql\",\n",
" \"modifier\": \"NOT_NULL\"\n",
" }\n",
" },\n",
" fragment=\"\"\"\n",
" id\n",
" stash_ids {\n",
" endpoint\n",
" stash_id\n",
" }\n",
" studio {\n",
" id\n",
" name\n",
" parent_studio {\n",
" id\n",
" name\n",
" }\n",
" }\n",
" \"\"\"\n",
")\n",
"\n",
"df_scenes_with_stashdb = pd.DataFrame(scenes_with_stashdb)\n",
"df_scenes_with_stashdb[\"studio_id\"] = df_scenes_with_stashdb[\"studio\"].apply(lambda x: x[\"id\"])\n",
"df_scenes_with_stashdb[\"studio_name\"] = df_scenes_with_stashdb[\"studio\"].apply(lambda x: x[\"name\"])\n",
"df_scenes_with_stashdb[\"parent_studio_id\"] = df_scenes_with_stashdb[\"studio\"].apply(lambda x: x[\"parent_studio\"][\"id\"] if x[\"parent_studio\"] else None)\n",
"df_scenes_with_stashdb[\"parent_studio_name\"] = df_scenes_with_stashdb[\"studio\"].apply(lambda x: x[\"parent_studio\"][\"name\"] if x[\"parent_studio\"] else None)\n",
"df_scenes_with_stashdb[\"is_tpdb_endpoint_missing\"] = df_scenes_with_stashdb[\"stash_ids\"].apply(is_tpdb_endpoint_missing)\n",
"\n",
"# Group by parent_studio_id and studio_id to count missing TPDB scenes\n",
"df_scenes_grouped_by_studio = df_scenes_with_stashdb[df_scenes_with_stashdb[\"is_tpdb_endpoint_missing\"]].groupby(\n",
" [\"parent_studio_id\", \"parent_studio_name\", \"studio_id\", \"studio_name\"]\n",
").size().reset_index(name=\"count_missing_tpdb\")\n",
"\n",
"# Calculate total scenes missing TPDB per parent studio\n",
"parent_totals = df_scenes_grouped_by_studio.groupby(\"parent_studio_id\")[\"count_missing_tpdb\"].sum().reset_index(name=\"parent_total_missing\")\n",
"\n",
"# Merge totals back into the grouped DataFrame\n",
"df_scenes_grouped_by_studio = df_scenes_grouped_by_studio.merge(parent_totals, on=\"parent_studio_id\")\n",
"\n",
"# Sort by parent studio total count and then by studio count\n",
"df_scenes_grouped_by_studio = df_scenes_grouped_by_studio.sort_values(by=[\"parent_total_missing\", \"count_missing_tpdb\"], ascending=[False, False])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Run this cell to find scenes of a single studio\n",
"studio_name = \"Meana Wolf\"\n",
"studios = stash.find_studios({ \n",
" \"name\": {\n",
" \"modifier\": \"EQUALS\",\n",
" \"value\": studio_name\n",
" } \n",
"})\n",
"\n",
"studio_ids = [studio[\"id\"] for studio in studios]\n",
"studio_ids"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load data from Stash\n",
"\n",
"stashapp_scenes = stash.find_scenes(\n",
" {\n",
" \"stash_id_endpoint\": {\n",
" \"endpoint\": \"https://theporndb.net/graphql\",\n",
" \"modifier\": \"IS_NULL\"\n",
" },\n",
" \"studios\": {\n",
" \"value\": studio_ids,\n",
" \"modifier\": \"INCLUDES\"\n",
" }\n",
" },\n",
" {\n",
" \"per_page\": 100,\n",
" \"sort\": \"date\"\n",
" },\n",
" fragment=\"\"\"\n",
" id\n",
" stash_ids {\n",
" endpoint\n",
" stash_id\n",
" }\n",
" title\n",
" date\n",
" urls\n",
" studio {\n",
" name\n",
" url\n",
" }\n",
" files {\n",
" fingerprints {\n",
" type\n",
" value\n",
" }\n",
" }\n",
" \"\"\"\n",
")\n",
"df_stashapp_scenes = pd.DataFrame(stashapp_scenes)\n",
"\n",
"def get_single_phash_from_files(files):\n",
" if len(files) != 1:\n",
" return None\n",
"\n",
" file = files[0]\n",
" if not \"fingerprints\" in file:\n",
" return None\n",
" \n",
" fingerprints = file[\"fingerprints\"]\n",
" if not isinstance(fingerprints, list):\n",
" return None\n",
" \n",
" # find a fingerprint with type phash\n",
" for fingerprint in fingerprints:\n",
" if fingerprint[\"type\"] == \"phash\":\n",
" return fingerprint[\"value\"]\n",
"\n",
" return None\n",
"\n",
"def get_stashdb_uuid(stash_ids):\n",
" for stash_id in stash_ids:\n",
" if stash_id[\"endpoint\"] == \"https://stashdb.org/graphql\":\n",
" return stash_id[\"stash_id\"]\n",
" return None\n",
"\n",
"df_stashapp_scenes\n",
"\n",
"# Filter dataframe to only include rows where the TPDB endpoint is missing\n",
"df_stashapp_scenes = df_stashapp_scenes[df_stashapp_scenes[\"stash_ids\"].apply(is_tpdb_endpoint_missing)]\n",
"df_stashapp_scenes[\"phash\"] = df_stashapp_scenes[\"files\"].apply(get_single_phash_from_files)\n",
"df_stashapp_scenes[\"stashdb_uuid\"] = df_stashapp_scenes[\"stash_ids\"].apply(get_stashdb_uuid)\n",
"df_stashapp_scenes[\"studio_name\"] = df_stashapp_scenes[\"studio\"].apply(lambda x: x[\"name\"])\n",
"\n",
"df_stashapp_scenes = df_stashapp_scenes.drop(columns=[\"studio\"])\n",
"\n",
"if df_stashapp_scenes.empty:\n",
" print(\"No scenes found with missing TPDB endpoint\")\n",
"else:\n",
" df_stashapp_scenes"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Match with TPDB\n",
"\n",
"import pandas as pd\n",
"import requests\n",
"\n",
"# Function to fetch scene data by phash and store directly in a list\n",
"def fetch_scene_data(df, headers):\n",
" # List to store the scene data for DataFrame creation\n",
" scenes_data_list = []\n",
"\n",
" # Loop through data frame rows\n",
" for index, row in df.iterrows():\n",
" id = row[\"id\"]\n",
" phash = row[\"phash\"]\n",
"\n",
" response = requests.get(\n",
" f\"https://api.theporndb.net/scenes?hash={phash}&hashType=PHASH\",\n",
" headers=headers,\n",
" )\n",
"\n",
" if response.status_code == 200:\n",
" response_json = response.json()\n",
"\n",
" if response_json:\n",
" scenes_data = response_json.get(\"data\", [])\n",
" if scenes_data:\n",
" scene_data_list = []\n",
" for scene_data in scenes_data:\n",
" scene_data_list.append({\n",
" \"tpdb_date\": scene_data.get(\"date\"),\n",
" \"tpdb_title\": scene_data.get(\"title\"),\n",
" \"tpdb_uuid\": scene_data.get(\"id\"),\n",
" \"tpdb_url\": scene_data.get(\"url\"),\n",
" \"tpdb_studio_id\": scene_data.get(\"site\", {}).get(\"uuid\"),\n",
" \"tpdb_studio_name\": scene_data.get(\"site\", {}).get(\"name\"),\n",
" \"tpdb_studio_url\": scene_data.get(\"site\", {}).get(\"url\"),\n",
" })\n",
" scenes_data_list.append({\n",
" \"id\": id,\n",
" \"stash_scene\": row.to_dict(), # Store the whole row as a dict\n",
" \"matches\": scene_data_list\n",
" })\n",
"\n",
" # Create a DataFrame from the list of scene data\n",
" return pd.DataFrame(scenes_data_list)\n",
"\n",
"# Fetch all scene data and store in a DataFrame\n",
"df_tpdb_scenes = fetch_scene_data(df_stashapp_scenes, tpdb_headers)\n",
"\n",
"if df_tpdb_scenes.empty:\n",
" print(\"No TPDB data found\")\n",
"else:\n",
" df_tpdb_scenes\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Analyzing the results\n",
"\n",
"import pandas as pd\n",
"from difflib import SequenceMatcher\n",
"\n",
"# Function to evaluate a match\n",
"def evaluate_match(stash_scene, match):\n",
" score = 0\n",
"\n",
" # Check date similarity\n",
" if pd.notna(match[\"tpdb_date\"]):\n",
" row_date = pd.to_datetime(stash_scene[\"date\"])\n",
" if abs(row_date - pd.to_datetime(match[\"tpdb_date\"])) < pd.Timedelta(days=1):\n",
" score += 1\n",
"\n",
" # Check title similarity\n",
" if pd.notna(match[\"tpdb_title\"]):\n",
" clean_title = stash_scene[\"title\"].strip().upper()\n",
" clean_tpdb_title = match[\"tpdb_title\"].strip().upper()\n",
" if clean_title == clean_tpdb_title:\n",
" score += 2 # Higher weight for exact match\n",
" elif SequenceMatcher(None, clean_title, clean_tpdb_title).ratio() > 0.9:\n",
" score += 1\n",
"\n",
" # Check studio similarity\n",
" if pd.notna(match[\"tpdb_studio_name\"]):\n",
" if stash_scene[\"studio_name\"] == match[\"tpdb_studio_name\"]:\n",
" score += 1\n",
"\n",
" # Check URL similarity\n",
" if pd.notna(match[\"tpdb_url\"]) and isinstance(match[\"tpdb_url\"], str):\n",
" clean_tpdb_url = match[\"tpdb_url\"].replace(\"www.\", \"\")\n",
" if stash_scene[\"urls\"]:\n",
" for url in stash_scene[\"urls\"]:\n",
" clean_url = url.replace(\"www.\", \"\")\n",
" if clean_url == clean_tpdb_url:\n",
" score += 1\n",
" break\n",
"\n",
" return score\n",
"\n",
"# Function to select the best match\n",
"def select_best_match(row):\n",
" best_match = None\n",
" highest_score = 0\n",
" stash_scene = row[\"stash_scene\"]\n",
"\n",
" for match in row[\"matches\"]:\n",
" score = evaluate_match(stash_scene, match)\n",
" if score > highest_score:\n",
" highest_score = score\n",
" best_match = match\n",
"\n",
" return best_match\n",
"\n",
"# Apply the selection function to get the best match and normalize it into a DataFrame\n",
"df_tpdb_scenes[\"best_match\"] = df_tpdb_scenes.apply(select_best_match, axis=1)\n",
"\n",
"# Flatten the best match columns into separate columns in the main DataFrame\n",
"best_match_df = pd.json_normalize(df_tpdb_scenes[\"best_match\"])\n",
"\n",
"# Merge the best match back into the original DataFrame\n",
"df_matched_scenes = pd.concat([df_stashapp_scenes, best_match_df], axis=1)\n",
"\n",
"\n",
"# Assume `df` and `scene_df` are already defined and contain the relevant data\n",
"# Assume `fetch_scene_data` is previously defined and executed\n",
"\n",
"# Function to evaluate a match\n",
"def evaluate_match(stash_scene, match):\n",
" score = 0\n",
"\n",
" # Check date similarity\n",
" if pd.notna(match[\"tpdb_date\"]):\n",
" row_date = pd.to_datetime(stash_scene[\"date\"])\n",
" if abs(row_date - pd.to_datetime(match[\"tpdb_date\"])) < pd.Timedelta(days=1):\n",
" score += 1\n",
"\n",
" # Check title similarity\n",
" if pd.notna(match[\"tpdb_title\"]):\n",
" clean_title = stash_scene[\"title\"].strip().upper()\n",
" clean_tpdb_title = match[\"tpdb_title\"].strip().upper()\n",
" if clean_title == clean_tpdb_title:\n",
" score += 2 # Higher weight for exact match\n",
" elif SequenceMatcher(None, clean_title, clean_tpdb_title).ratio() > 0.9:\n",
" score += 1\n",
"\n",
" # Check studio similarity\n",
" if pd.notna(match[\"tpdb_studio_name\"]):\n",
" if stash_scene[\"studio_name\"] == match[\"tpdb_studio_name\"]:\n",
" score += 1\n",
"\n",
" # Check URL similarity\n",
" if pd.notna(match[\"tpdb_url\"]) and isinstance(match[\"tpdb_url\"], str):\n",
" clean_tpdb_url = match[\"tpdb_url\"].replace(\"www.\", \"\")\n",
" if stash_scene[\"urls\"]:\n",
" for url in stash_scene[\"urls\"]:\n",
" clean_url = url.replace(\"www.\", \"\")\n",
" if clean_url == clean_tpdb_url:\n",
" score += 1\n",
" break\n",
"\n",
" return score\n",
"\n",
"# Function to select the best match\n",
"def select_best_match(row):\n",
" best_match = None\n",
" highest_score = 0\n",
" stash_scene = row[\"stash_scene\"]\n",
"\n",
" for match in row[\"matches\"]:\n",
" score = evaluate_match(stash_scene, match)\n",
" if score > highest_score:\n",
" highest_score = score\n",
" best_match = match\n",
"\n",
" return best_match\n",
"\n",
"# Apply the selection function to get the best match\n",
"df_tpdb_scenes[\"best_match\"] = df_tpdb_scenes.apply(select_best_match, axis=1)\n",
"\n",
"# Normalize the best match into separate columns\n",
"best_match_df = pd.json_normalize(df_tpdb_scenes[\"best_match\"])\n",
"\n",
"# Combine the best matches with the original scene_df\n",
"df_tpdb_scenes = pd.concat([df_tpdb_scenes.drop(columns=[\"best_match\", \"matches\"]), best_match_df], axis=1)\n",
"\n",
"# Merge with the original DataFrame on the 'id' column\n",
"df_matched_scenes = pd.merge(df_stashapp_scenes, df_tpdb_scenes, on=\"id\", how=\"inner\")\n",
"\n",
"# Clean up columns after merging if needed\n",
"# Remove the stash_scene column after merging if it's no longer needed\n",
"df_matched_scenes = df_matched_scenes.drop(columns=[\"stash_scene\"])\n",
"\n",
"# Function to check date and title similarity (unchanged)\n",
"def analyze_scene(row):\n",
" # Initialize matching results\n",
" date_match = False\n",
" exact_title_match = False\n",
" near_title_match = False\n",
" studio_match = False\n",
" url_match = False\n",
"\n",
" if pd.notna(row[\"tpdb_date\"]):\n",
" # Check date similarity\n",
" row_date = pd.to_datetime(row[\"date\"])\n",
" if abs(row_date - pd.to_datetime(row[\"tpdb_date\"])) < pd.Timedelta(days=1):\n",
" date_match = True\n",
"\n",
" if pd.notna(row[\"tpdb_title\"]):\n",
" # Check title similarity\n",
" clean_title = row[\"title\"].strip().upper()\n",
" clean_tpdb_title = row[\"tpdb_title\"].strip().upper()\n",
" exact_title_match = clean_title == clean_tpdb_title\n",
" near_title_match = SequenceMatcher(None, clean_title, clean_tpdb_title).ratio() > 0.9\n",
" \n",
" if pd.notna(row[\"tpdb_studio_name\"]):\n",
" # Check studio similarity\n",
" if row[\"studio_name\"] == row[\"tpdb_studio_name\"]:\n",
" studio_match = True\n",
" \n",
" if pd.notna(row[\"tpdb_url\"]):\n",
" # Check URL is string\n",
" if not isinstance(row[\"tpdb_url\"], str):\n",
" url_match = False\n",
" print(f\"URL is not a string: {row['id']} {row['tpdb_uuid']} {row['tpdb_url']}\")\n",
" else:\n",
" clean_tpdb_url = row[\"tpdb_url\"].replace(\"www.\", \"\")\n",
"\n",
" if row[\"urls\"]:\n",
" for url in row[\"urls\"]:\n",
" clean_url = url.replace(\"www.\", \"\")\n",
" if clean_url == clean_tpdb_url:\n",
" url_match = True\n",
" \n",
"\n",
" return pd.Series({\n",
" \"date_match\": date_match,\n",
" \"exact_title_match\": exact_title_match,\n",
" \"near_title_match\": near_title_match,\n",
" \"studio_match\": studio_match,\n",
" \"url_match\": url_match,\n",
" \"all_match\": all([date_match, exact_title_match, near_title_match, studio_match, url_match])\n",
" })\n",
"\n",
"# Apply the analysis and add results to new columns in the merged DataFrame\n",
"df_matched_scenes[[\"date_match\", \"exact_title_match\", \"near_title_match\", \"studio_match\", \"url_match\", \"all_match\"]] = df_matched_scenes.apply(analyze_scene, axis=1)\n",
"\n",
"# Define the new column order\n",
"new_column_order = [\n",
" 'id', \n",
" \n",
" 'all_match', 'date_match', 'exact_title_match', 'near_title_match', 'studio_match', 'url_match',\n",
" \n",
" 'stashdb_uuid', 'tpdb_uuid',\n",
"\n",
" 'title', 'tpdb_title',\n",
" 'date', 'tpdb_date',\n",
" 'studio_name', 'tpdb_studio_name', 'tpdb_studio_id', 'tpdb_studio_url',\n",
" 'urls', 'tpdb_url',\n",
" 'files',\n",
" 'phash',\n",
"\n",
" 'stash_ids',\n",
"]\n",
"\n",
"# Reorder the DataFrame columns\n",
"df_matched_scenes = df_matched_scenes[new_column_order]\n",
"df_matched_scenes = df_matched_scenes[df_matched_scenes[\"tpdb_uuid\"].notna()]\n",
"\n",
"df_matched_scenes\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Run this cell to save TPDB UUID for those where all identifiers match\n",
"save_tpdb_uuid(df_matched_scenes[df_matched_scenes[\"all_match\"]])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Example: Run this cell to save TPDB UUID for those where all date and studio name match and title is a near match.\n",
"# Tweak the conditions as needed.\n",
"df_partially_matched = df_matched_scenes[df_matched_scenes[\"date_match\"] & df_matched_scenes[\"near_title_match\"] & df_matched_scenes[\"studio_match\"]]\n",
"save_tpdb_uuid(df_partially_matched)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Run this cell if you have manually verified that all rows match. This is useful for example in case StashDB and TPDB have different\n",
"# URLs for the same scene, but all other identifiers match.\n",
"save_tpdb_uuid(df_matched_scenes)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "pandas",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.4"
}
},
"nbformat": 4,
"nbformat_minor": 2
}