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)