generate_info_dic.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import warnings\n",
    "from easydict import EasyDict as edict\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from utils.data_preprocess_utils import get_all_files, get_config\n",
    "from easydict import EasyDict as edict\n",
    "from tqdm import tqdm\n",
    "import dill\n",
    "cfg = edict(get_config('./cfg/data_preprocess_cfg.yaml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 184/184 [06:05<00:00,  1.99s/it]\n"
     ]
    }
   ],
   "source": [
    "from os import path\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import dill\n",
    "\n",
    "def process_file(file):\n",
    "    if file.endswith('.csv'):\n",
    "        df = pd.read_csv(file, usecols=['MMSI'])\n",
    "        return df['MMSI'].unique()\n",
    "    else:\n",
    "        return []\n",
    "\n",
    "files = get_all_files(cfg.in_dir)\n",
    "if path.exists(cfg.mmsi_info_dir):\n",
    "    with open(cfg.mmsi_info_dir, 'rb') as f:\n",
    "        MMSIs = dill.load(f)\n",
    "else:\n",
    "    MMSIs = []\n",
    "\n",
    "with ThreadPoolExecutor(max_workers=8) as executor:\n",
    "    results = list(tqdm(executor.map(process_file, files), total=len(files)))\n",
    "\n",
    "for result in results:\n",
    "    MMSIs.extend(result)\n",
    "\n",
    "MMSIs = list(set(MMSIs))\n",
    "\n",
    "with open(cfg.mmsi_info_dir, 'wb') as f:\n",
    "    dill.dump(MMSIs, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bs4 import BeautifulSoup\n",
    "import json\n",
    "import requests\n",
    "def get_vessel_info(mmsi:str):\n",
    "    '''\n",
    "    Parameters:\n",
    "        mmsi: MMSI of the vessel\n",
    "\n",
    "    Returns:\n",
    "        vessel_type: Pramary type, String,\n",
    "        vessel_subtype: Pramary scendary type, String\n",
    "        vessel_length: length (dm), float, Note: 10dm = 1m\n",
    "        vessel_width: vessel width (dm), float\n",
    "\n",
    "    Typical usage example:\n",
    "        type_cal = ['Sailing','Pleasure' ,'Cargo','Fishing','Passenger','Tanker','Tug','SAR','HSC','Dredging','Military']\n",
    "\n",
    "        mmsi = '538003672'\n",
    "\n",
    "        vessel_type, vessel_subtype, vessel_length, vessel_width = get_vessel_info(type_cal,mmsi)\n",
    "        \n",
    "    '''\n",
    "    type_cal = ['Sailing','Pleasure' ,'Cargo','Fishing','Passenger','Tanker','Tug','SAR','HSC','Dredging','Military']\n",
    "\n",
    "    vessel_type = 'Other'\n",
    "    vessel_subtype = ''\n",
    "    vessel_length = ''\n",
    "    vessel_width = ''\n",
    "    headers = {\n",
    "        \"User-Agent\": \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36\",\n",
    "        \"Referer\": \"https://www.marinetraffic.com/\",\n",
    "        \"Accept-Language\": \"en-US,en;q=0.9\",\n",
    "    }\n",
    "\n",
    "    data = {'mmsi':mmsi}\n",
    "    url = f'https://www.marinetraffic.com/en/ais/details/ships/mmsi:{mmsi}'\n",
    "    \n",
    "    try :\n",
    "        response = requests.get(url=url, headers=headers)\n",
    "    except requests.exceptions.ConnectionError:\n",
    "        print('ConnectionError')\n",
    "        return vessel_type, vessel_subtype, vessel_length, vessel_width\n",
    "    if response.status_code == 400:\n",
    "        print('Invalid MMSI')\n",
    "    if response.status_code == 200:\n",
    "        soup = BeautifulSoup(response.content, 'html.parser')\n",
    "        content = soup.find_all('meta', {'content': lambda x: x and x.startswith('vessel ')})\n",
    "        for tag in content:\n",
    "            if tag.get('data-react-helmet') == 'true' and tag.get('name') == 'keywords':\n",
    "                keywords = tag.get('content')\n",
    "                vessel_type = keywords.split(',')[3][1:]\n",
    "                vessel_type = vessel_type.split(' ')[0]\n",
    "                vessel_subtype = keywords.split(',')[2][13:]\n",
    "                if vessel_type not in type_cal:\n",
    "                    vessel_type = 'Other'\n",
    "                break\n",
    "\n",
    "    try: \n",
    "        response = requests.post(\"http://www.shipfinder.com/ship/GetShip\", data=data)\n",
    "    except requests.exceptions.ConnectionError:\n",
    "        print('ConnectionError')\n",
    "        return vessel_type, vessel_subtype, vessel_length, vessel_width\n",
    "    if response.status_code == 200:\n",
    "        if json.loads(response.text)['data'] == []:\n",
    "            return vessel_type, vessel_subtype, vessel_length, vessel_width\n",
    "        result = json.loads(response.text)['data'][0]\n",
    "        vessel_length = result['length']\n",
    "        vessel_width = result['width']\n",
    "        vessel_length = float(vessel_length)\n",
    "        vessel_width = float(vessel_width)\n",
    "    return vessel_type, vessel_subtype, vessel_length, vessel_width\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 110777/110777 [06:04<00:00, 304.18it/s]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import concurrent.futures\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def process_mmsi(mmsi):\n",
    "    mmsi_str = str(mmsi)\n",
    "    if not (pattern.match(mmsi_str) and (len(mmsi_str) == 9)):\n",
    "        return None\n",
    "    if mmsi_str not in ships_info_dict.keys():\n",
    "        ships_info_dict[mmsi_str] = {}\n",
    "        vessel_type, vessel_subtype, vessel_length, vessel_width = get_vessel_info(mmsi_str)\n",
    "        ships_info_dict[mmsi_str]['vessel_type'] = vessel_type\n",
    "        ships_info_dict[mmsi_str]['vessel_subtype'] = vessel_subtype\n",
    "        ships_info_dict[mmsi_str]['vessel_length'] = vessel_length\n",
    "        ships_info_dict[mmsi_str]['vessel_width'] = vessel_width\n",
    "    return mmsi_str\n",
    "\n",
    "if not os.path.exists(cfg.vessels_info_dic_dir):\n",
    "    ships_info_dict = {}\n",
    "else:\n",
    "    with open(cfg.vessels_info_dic_dir) as json_file:\n",
    "        ships_info_dict = json.load(json_file)\n",
    "\n",
    "pattern = re.compile(r'^[2-7]')\n",
    "with open(cfg.mmsi_info_dir, 'rb') as f:\n",
    "    MMSIs = dill.load(f)\n",
    "\n",
    "# 使用ThreadPoolExecutor进行多线程处理\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    results = list(tqdm(executor.map(process_mmsi, MMSIs), total=len(MMSIs)))\n",
    "\n",
    "# 将处理完的结果保存到文件\n",
    "with open(cfg.vessels_info_dic_dir, 'w') as fp:\n",
    "    json.dump(ships_info_dict, fp)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rapids-23.10",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}