forked from 22BobbyW/Sinkers
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Image_Processor.py
99 lines (75 loc) · 3.22 KB
/
Image_Processor.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jul 22 11:30:37 2021
@author: BWSI AUV Challenge Instructional Staff
"""
### JRE: for simulation only!
import sys
import pathlib
import datetime
import time
import cv2
import numpy as np
import os
if os.uname().nodename == 'sharkpi' or os.uname().nodename == 'auvpi':
import picamera
# For simulations
from BWSI_BuoyField import BuoyField
from BWSI_Sensor import BWSI_Camera
from cam_util import detect_buoys
class ImageProcessor():
def __init__(self, camera='PICAM', log_dir='./'):
self.__camera_type = camera.upper()
if self.__camera_type == 'SIM':
self.__camera = BWSI_Camera(max_angle=31.1, visibility=30)
self.__simField = None
else:
self.__camera = picamera.PiCamera()
self.__camera.resolution = (640, 480)
self.__camera.framerate = 24
time.sleep(2) # camera warmup time
self.__image = np.empty((480*640*3,), dtype=np.uint8)
# create my save directory
self.__image_dir = pathlib.Path(log_dir, 'frames')
self.__image_dir.mkdir(parents=True, exist_ok=True)
# ------------------------------------------------------------------------ #
# Run an iteration of the image processor.
# The sim version needs the auv state ot generate simulated imagery
# the PICAM does not need any auv_state input
# ------------------------------------------------------------------------ #
def run(self, auv_state=None):
red = None
green = None
if auv_state['heading'] is not None:
if (self.__camera_type == 'SIM'):
# if it's the first time through, configure the buoy field
if self.__simField is None:
self.__simField = BuoyField(auv_state['datum'])
config = {'nGates': 5,
'gate_spacing': 5,
'gate_width': 2,
'style': 'pool_1',
'max_offset': 5,
'heading': 45}
self.__simField.configure(config)
# synthesize an image
image = self.__camera.get_frame(auv_state['position'], auv_state['heading'], self.__simField)
elif self.__camera_type == 'PICAM':
try:
self.__camera.capture(self.__image, 'bgr')
except:
# restart the camera
# self.__camera = picamera.PiCamera()
self.__camera.resolution = (640, 480)
self.__camera.framerate = 24
time.sleep(2) # camera warmup time
image = self.__image.reshape((480, 640, 3))
else:
print(f"Unknown camera type: {self.__camera_type}")
sys.exit(-10)
# log the image
fn = self.__image_dir / f"frame_{int(datetime.datetime.utcnow().timestamp())}.jpg"
cv2.imwrite(str(fn), image)
green, red = detect_buoys(image)
return green, red