{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "238f7b3e-852a-46d2-a5e6-bd86b537e59d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ultralytics gdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e680a9-2c32-4702-9be7-aadd33c3bff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import nltk\n",
    "nltk.download(['punkt', 'averaged_perceptron_tagger', 'averaged_perceptron_tagger_eng', 'wordnet', 'omw-1.4'])\n",
    "from nltk.tokenize import word_tokenize\n",
    "from nltk.tag import pos_tag\n",
    "from nltk.stem import WordNetLemmatizer\n",
    "from ultralytics import YOLO\n",
    "import gdown\n",
    "import numpy as np\n",
    "drive_link = \"update_only_this_line\"\n",
    "prompt_dataset = \"https://docs.google.com/document/d/1VJHTKDMPapyLbACAIOWvzRdrUhV70y03Mbyr_y--AN0/edit?usp=sharing\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45192853-a99d-4185-9359-a413095681f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_yolo_file():\n",
    "    from ultralytics import YOLO\n",
    "    model = YOLO('yolov8n.pt')\n",
    "    return 'yolov8n.pt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbc203f3-a618-4ced-86f5-10cea2f6d73a",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolo_model = download_yolo_file()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22ba5e05-aad8-47ef-b6a6-8addbe5f95ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_results_from_drive(drive_link):\n",
    "    gdown.download_folder(drive_link, output='kaggle_resources', quiet=False, use_cookies=False)\n",
    "    csv_data = pd.read_csv('kaggle_resources/results.csv')\n",
    "    return csv_data[['run_id', 'prompt', 'filenames']].rename(columns={'filenames': 'generated_images'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b282fe25-b2f0-48f9-8f55-ac5c2e348f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_data = parse_results_from_drive(drive_link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7160670-b3e4-4abb-9901-dcb70acc9a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c72e0c2-3b29-41c2-8ec1-f2f83a9af39a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_object_detection(yolo_file, generated_df):\n",
    "    model = YOLO(yolo_file)\n",
    "    def detect_objects(img):\n",
    "        try:\n",
    "            detections = model(os.path.join('kaggle_resources', img))[0]\n",
    "            return [model.names[int(cls)] for cls in detections.boxes.cls] if detections.boxes is not None else []\n",
    "        except:\n",
    "            return []\n",
    "    generated_df['predicted_objects'] = [detect_objects(img) for img in generated_df['generated_images']]\n",
    "    return generated_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab1da815-314e-4607-a118-db0b302b4356",
   "metadata": {},
   "outputs": [],
   "source": [
    "submissions_df = apply_object_detection(yolo_model, generated_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f82368-601e-4b41-a8f8-fbe4289234db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_ground_truth(drive_link):\n",
    "    import nltk\n",
    "    from nltk.tokenize import word_tokenize\n",
    "    from nltk.tag import pos_tag\n",
    "    from nltk.stem import WordNetLemmatizer\n",
    "    \n",
    "    lemmatizer = WordNetLemmatizer()\n",
    "    doc_id = drive_link.split('/d/')[1].split('/')[0]\n",
    "    export_url = f\"https://docs.google.com/document/d/{doc_id}/export?format=txt\"\n",
    "    gdown.download(export_url, 'temp_prompts.txt', quiet=False)\n",
    "    with open('temp_prompts.txt', 'r') as f:\n",
    "        prompts = f.readlines()\n",
    "    \n",
    "    results = []\n",
    "    for i, prompt in enumerate(prompts):\n",
    "        tokens = word_tokenize(prompt.lower().strip())\n",
    "        pos_tags = pos_tag(tokens)\n",
    "        nouns = [lemmatizer.lemmatize(word) for word, pos in pos_tags if pos.startswith('NN')]\n",
    "        results.append({'ID': i+1, 'prompt_id': i, 'prompt': prompt.strip(), 'ground_truth': nouns, 'Usage': 'Public'})\n",
    "    \n",
    "    return pd.DataFrame(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b09dbcde-a3ad-4a65-8ca7-6a5af7d029b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "solutions_df = extract_ground_truth(prompt_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "594e4fb0-12ca-4c07-877a-178d98a5e920",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_id_to_submission(submission_df, solution_df):\n",
    "    merged = solution_df[['ID', 'prompt']].merge(submission_df, on='prompt', how='left')\n",
    "    merged['predicted_objects'] = merged['predicted_objects'].fillna('').apply(lambda x: x if isinstance(x, list) else [])\n",
    "    merged[['run_id', 'generated_images']] = merged[['run_id', 'generated_images']].fillna(\"missing_information\")\n",
    "    return merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f6f0bd2-172d-4ace-9569-51bbde0b7060",
   "metadata": {},
   "outputs": [],
   "source": [
    "submissions_df = add_id_to_submission(submissions_df, solutions_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a1b3f4b-322b-4f3a-8fac-b4bd9beb3dbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(submissions_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73e4611c-4401-4ede-88bb-7ce2ff8c7702",
   "metadata": {},
   "outputs": [],
   "source": [
    "def score(solution: pd.DataFrame, submission: pd.DataFrame) -> float:\n",
    "    \"\"\"\n",
    "    Calculate the average F1 score between predicted and ground truth objects using proper precision/recall.\n",
    "    \n",
    "    Args:\n",
    "        solution: DataFrame with 'ID' and 'ground_truth' columns\n",
    "        submission: DataFrame with 'ID' and 'predicted_objects' columns\n",
    "    \n",
    "    Returns:\n",
    "        float: Average F1 score across all prompts using TP, FP, FN calculations\n",
    "    \"\"\"\n",
    "    import ast\n",
    "    merged = solution.merge(submission, on='ID', how='left')\n",
    "    merged['predicted_objects'] = merged['predicted_objects'].fillna('[]').apply(lambda x: ast.literal_eval(x) if isinstance(x, str) else x)\n",
    "    merged['ground_truth'] = merged['ground_truth'].apply(lambda x: ast.literal_eval(x) if isinstance(x, str) else x)\n",
    "    \n",
    "    f1_scores = []\n",
    "    for pred, truth in zip(merged['predicted_objects'], merged['ground_truth']):\n",
    "        pred_set, truth_set = set(pred), set(truth)\n",
    "        tp = len(pred_set & truth_set)\n",
    "        fp = len(pred_set - truth_set)\n",
    "        fn = len(truth_set - pred_set)\n",
    "        precision = tp / (tp + fp) if (tp + fp) > 0 else 0\n",
    "        recall = tp / (tp + fn) if (tp + fn) > 0 else 0\n",
    "        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0\n",
    "        f1_scores.append(f1)\n",
    "    return np.mean(f1_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08ffcd47-40ee-49bc-afc6-ad60500985dc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "f1_score = score(solutions_df, submissions_df)\n",
    "print(f1_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a5fbca8-faac-469a-8bfe-f2419deabe4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "submissions_df.to_csv(\"submission.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a0c78f2-0a95-423a-853c-659d5a149b61",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}