Initial commit

.env

@@ -0,0 +1 @@
+SERVER_IP=localhost #192.168.174.39

.gitignore

@@ -0,0 +1,2 @@
+venv/
+best.pt

README.md

@@ -0,0 +1,37 @@
+# Object detection based target following
+Authors: Murad Bashirov, Mehdi Aghakishiyev, Khadija Rajabova
+
+## Getting Started
+### Requirements
+**Python** >= 3.8 and **PyTorch**>=1.7 (see https://pytorch.org/get-started/locally)
+1. Clone the repository
+```shell
+git clone https://github.com/m-spitfire/coe202-project
+```
+### For server
+1. Download the `best.pt` pre-trained model from the [Releases](https://github.com/m-spitfire/coe202-project/releases) tab.
+2. Install the requirements for `yolov5`:
+```shell
+pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
+```
+3. Install the requirements for server:
+```shell
+pip install -r requirements-server.txt
+```
+4. Get your local IP address and set it in the [`.env`](./.env) file
+5. Run the server:
+```shell
+python server.py
+```
+### For client
+1. Install the dependencies for client:
+```shell
+pip install -r requirements-client.txt
+```
+2. Run the client
+```shell
+python client.py
+```
+
+## Acknowledgements
+* [WelkinU/yolov5-fastapi-demo](https://github.com/WelkinU/yolov5-fastapi-demo) for basic fastapi server-client example

client.py

@@ -0,0 +1,39 @@
+import cv2
+import requests
+from dotenv import load_dotenv
+
+import time
+import os
+
+load_dotenv()
+ip = os.environ.get("SERVER_IP")
+
+cap = cv2.VideoCapture(0)
+cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
+cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 640)
+cap.set(cv2.CAP_PROP_FPS, 36)
+
+
+def infer(i):
+ ret, img = cap.read()
+ if not ret:
+ return
+
+ img = cv2.rotate(img, cv2.ROTATE_180)
+ retv, buf = cv2.imencode(".jpg", img)
+
+ result = requests.post(f"http://{ip}:8000", files={'file': buf}).json()
+
+ try:
+ if result[0][0]['confidence'] > 0.7:
+ cv2.imwrite(f"save{i}.jpg", img)
+ print(i, result)
+ except:
+ pass
+
+
+i = 0
+while True:
+ infer(i)
+ time.sleep(1)
+ i += 1

requirements-client.txt

@@ -0,0 +1,3 @@
+opencv-python
+requests
+python-dotenv

requirements-server.txt

@@ -0,0 +1,4 @@
+python-dotenv
+fastapi
+python-multipart
+uvicorn

server.py

@@ -0,0 +1,52 @@
+from fastapi import FastAPI, Request, Form, File, UploadFile, HTTPException
+from fastapi.responses import HTMLResponse
+from dotenv import load_dotenv
+
+from PIL import Image
+from io import BytesIO
+import os
+
+import torch
+
+load_dotenv()
+ip = os.environ.get("SERVER_IP") or "localhost"
+
+app = FastAPI()
+
+@app.get("/")
+async def home(request: Request):
+
+ return HTTPException(status_code=405)
+
+
+@app.post("/")
+async def detect(file: UploadFile = File(...)):
+
+ model = torch.hub.load('ultralytics/yolov5', 'custom', './best.pt')
+
+ results = model(Image.open(BytesIO(await file.read())))
+
+ json_results = results_to_json(results,model)
+ return json_results
+
+
+def results_to_json(results, model):
+ return [
+ [
+ {
+ "class": int(pred[5]),
+ "class_name": model.model.names[int(pred[5])],
+ "bbox": [int(x) for x in pred[:4].tolist()],
+ "confidence": float(pred[4]),
+ }
+ for pred in result
+ ]
+ for result in results.xyxy
+ ]
+
+
+if __name__ == '__main__':
+ import uvicorn
+
+ app_str = 'server:app'
+ uvicorn.run(app_str, host=ip, port=8000, reload=True, workers=1)