widzenie-komputerowe-projekt/rybki.ipynb

123 lines
3.6 KiB
Plaintext
Raw Normal View History

2023-01-30 17:00:59 +01:00
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import cv2\n",
"import matplotlib.pyplot as plt\n",
"import keras\n",
"import numpy as np\n",
"import threading"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"model = keras.models.load_model('./model')"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"cap = cv2.VideoCapture(\"rybki.mp4\")\n",
"cap.set(cv2.CAP_PROP_FPS, 60)\n",
"\n",
"class_names=sorted(['Fish', \"Jellyfish\", 'Penguin', 'Puffin', 'Shark', 'Starfish', 'Stingray'])\n",
"object_detector = cv2.createBackgroundSubtractorMOG2(history=100, varThreshold=50)\n",
"\n",
"\n",
"width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) \n",
"height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
"fps = cap.get(cv2.CAP_PROP_FPS)\n",
"out = cv2.VideoWriter('track_fish.avi', cv2.VideoWriter_fourcc(*'MJPG'), fps, (width, height))\n",
"\n",
"while True:\n",
" ret, frame = cap.read()\n",
" if(frame is not None):\n",
" roi = frame[100: 900,330:1900]\n",
" mask = object_detector.apply(roi)\n",
" _, mask = cv2.threshold(mask,254,255, cv2.THRESH_BINARY)\n",
" conturs, _ =cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)\n",
"\n",
" images = []\n",
"\n",
" for cnt in conturs:\n",
" area = cv2.contourArea(cnt)\n",
" if area > 200:\n",
" #cv2.drawContours(roi,[cnt],-1,(0,255,0),2)\n",
" x,y,w,h = cv2.boundingRect(cnt)\n",
" rectangle = cv2.rectangle(roi,(x,y),(x+w,y+h),(0,255,0),3)\n",
" image_to_predict = roi[y:y+h,x:x+w]\n",
" image_to_predict = cv2.resize(image_to_predict,(227,227))\n",
" images.append((x,y,rectangle,np.expand_dims(image_to_predict,axis=0)))\n",
" \n",
" # pred = model.predict(np.expand_dims(image_to_predict, axis=0))\n",
" # label = class_names[np.argmax(pred)]\n",
" if images:\n",
" pred = model.predict(np.vstack([image[3] for image in images]))\n",
" labels = [class_names[np.argmax(pre)] for pre in pred]\n",
" for i,image in enumerate(images):\n",
" cv2.putText(image[2], labels[i], (image[0], image[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 1)\n",
" roi = cv2.resize(roi, (960, 540)) \n",
" cv2.imshow(\"roi\", roi)\n",
"\n",
" key = cv2.waitKey(30)\n",
" if key == 27:\n",
" break\n",
"\n",
" #out.write(frame)\n",
" else:\n",
" break\n",
"\n",
"\n",
"#out.release()\n",
"cap.release()\n",
"cv2.destroyAllWindows()"
]
},
{
"cell_type": "code",
"execution_count": 59,
"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.9.2"
},
"orig_nbformat": 4,
"vscode": {
"interpreter": {
"hash": "393784674bcf6e74f2fc9b1b5fb3713f9bd5fc6f8172c594e5cfa8e8c12849bc"
}
}
},
"nbformat": 4,
"nbformat_minor": 2
}