competition/科大讯飞AI开发者大赛2021/科大讯飞股份有限公司_猪只盘点挑战赛.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5516f498",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from imageai.Detection import ObjectDetection\n",
    "\n",
    "# imageai说明请查看官网： https://github.com/OlafenwaMoses/ImageAI/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4ccc8ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "execution_path = os.getcwd()\n",
    "\n",
    "detector = ObjectDetection()\n",
    "detector.setModelTypeAsRetinaNet()\n",
    "detector.setModelPath( os.path.join(execution_path , \"resnet50_coco_best_v2.1.0.h5\")) # 需要提取从官网下载h5文件\n",
    "detector.loadModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3ae4b045",
   "metadata": {},
   "outputs": [],
   "source": [
    "num = []\n",
    "pigs = []\n",
    "pic = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b1e853f8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 注意，这里的路径必须是全英文路径\n",
    "for file in os.listdir('E:\\\\pig\\\\test'):\n",
    "    detections = detector.detectObjectsFromImage(input_image=os.path.join(execution_path , \"E:\\\\pig\\\\test\\\\\"+file), output_image_path=os.path.join(execution_path , \"E:\\\\pig\\\\test2\\\\\"+'2a'+file), minimum_percentage_probability=30)\n",
    "    num.append(len(detections))\n",
    "    pig = []\n",
    "    pic.append(file)\n",
    "    print(len(pic))\n",
    "    for eachObject in detections: \n",
    "        pig.append([eachObject[\"percentage_probability\"]]+eachObject[\"box_points\"])\n",
    "#         print(eachObject[\"name\"] , \" : \", eachObject[\"percentage_probability\"], \" : \", eachObject[\"box_points\"] )\n",
    "#         print(\"--------------------------------\")\n",
    "    pigs.append(pig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f16b5c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 写文件pig_count.txt\n",
    "with open('E:\\\\pig\\\\det_files\\\\pig_count.txt','w',encoding='utf-8') as f:\n",
    "    [f.write('{0} {1}\\n'.format(key, value)) for key,value in zip(pic,num)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a40234e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 写文件det_results\n",
    "for i in range(len(num)):\n",
    "    with open('E:\\\\pig\\\\det_files\\\\det_results\\\\'+pic[i].split('.')[0]+'.txt','w',encoding='utf-8') as f:\n",
    "        [f.write('pig {0} {1} {2} {3} {4}\\n'.format(min(0.01*value[0]+0.4,0.99),(value[3]+value[1])/2,(value[2]+value[4])/2,(value[3]-value[1]),(value[4]-value[2]))) for value in pigs[i]]\n",
    "        # ****需要注意，最终提交的格式是x_ccenter,y_center,w,h和训练集里面的box格式不同。这个当时提交踩了很多坑"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68675b7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}