ppbbww.py (6680B)
1 import argparse 2 import asyncio 3 import json 4 import logging 5 import random 6 import tempfile 7 import time 8 import os 9 import sys 10 11 from datetime import datetime 12 from slack_sdk.web.async_client import AsyncWebClient 13 from PIL import Image, ImageDraw, ImageFont 14 15 from .archive import Archive 16 from .stream_sampler import StreamSampler 17 from .object_detector import ObjectDetector 18 19 20 async def sample_stream(frames_q, sampler, min_delay, max_delay, day_start, day_end): 21 fail_count, max_failures = 0, 5 22 while True: 23 try: 24 backoff = (2 ** fail_count) - 1 + random.uniform(0, 2*fail_count) 25 await asyncio.sleep(backoff) 26 27 frames = await sampler.get_recent_frames() 28 await frames_q.put(frames[0].path) 29 for frame in frames[1:]: # only keep 1 frame 30 os.remove(frame) 31 fail_count = 0 # this request succeeded, reset fail counter. 32 except Exception as ex: 33 logging.error(f"[sample_stream]: Failed to get_recent_frames: {ex}") 34 fail_count += 1 35 #if fail_count >= max_failures: 36 # logging.warning("[sample_stream]: Too many download failures. Killing program!") 37 # sys.exit(1) 38 39 if day_start <= datetime.now().time() <= day_end: 40 await asyncio.sleep(random.randint(min_delay, max_delay)) 41 else: 42 logging.warning(f"[sample_stream]: Day has ended!") 43 delay = 86400 - (datetime.combine(datetime.today(), day_end) - datetime.combine(datetime.today(), day_start)).total_seconds() 44 await asyncio.sleep(delay) 45 logging.warning(f"[sample_stream]: Day has begun!") 46 47 48 async def find_matches(frames_q, archive_q, announce_q): 49 detector = ObjectDetector(thresh=0.80) 50 while True: 51 # Get filenames from the queue and run object detector. 52 frame_file = await frames_q.get() 53 image = Image.open(frame_file) 54 matches = list(filter_matches(detector.find(image))) 55 if len(matches) == 0: 56 os.remove(frame_file) 57 continue 58 # Archive (and announce) the frame and matches. 59 await archive_q.put((frame_file, matches)) 60 await announce_q.put((frame_file, matches)) 61 62 63 def filter_matches(matches): 64 matches = list(matches) 65 num_boxes = len(matches) 66 for label, score, box in matches: 67 box = list(box) 68 # Skip matches at the bottom of the image area (rocks). 69 if box[1] > 600: 70 continue 71 # Skip common but uninteresting cases based on size and label. 72 box_size = (box[2] - box[0]) * (box[3] - box[1]) 73 if num_boxes < 3 and box_size < 500: 74 continue 75 if num_boxes < 3 and label == "boat" and box_size < 3000: 76 continue 77 if num_boxes < 5 and label == "bird" and box_size < 1000: 78 continue 79 # Return everything else. 80 yield (label, score, box) 81 82 83 async def archive_matches(archive_q, archive): 84 while True: 85 file, matches = await archive_q.get() 86 logging.warning(f"[archive_matches] Archiving {file}: {matches}") 87 ts = file.split("/")[-1].split("-")[0] 88 for label, score, box in matches: 89 archive.add_match(ts=ts, filename=file, label=label, score=score, box=box) 90 91 92 async def announce_matches(announce_q, client): 93 last_announce_ts, thread_ts = 0, None 94 mkdate = lambda ts: datetime.fromtimestamp(ts).date() 95 while True: 96 file, matches = await announce_q.get() 97 last_date, now_date = mkdate(last_announce_ts), mkdate(time.time()) 98 logging.warning(f"[announce_matches] last_date={last_date} now_date={now_date} thread_ts={thread_ts}") 99 thread_ts = await post_match( 100 client, file, thread_ts if last_date == now_date else None 101 ) 102 logging.warning(f"[annnounce_matches] Posted match at thread_ts {thread_ts}") 103 last_announce_ts = time.time() 104 105 106 async def post_match(client, frame_file, thread_ts): 107 chan_name, chan_id = "#boatwatch", "C06LMBRNV8U" 108 frame_ts = frame_file.split("/")[-1].split("-")[0] 109 with open(frame_file, "rb") as f: 110 try: 111 res = await client.files_upload_v2( 112 channel=chan_id, 113 thread_ts=thread_ts, 114 file=f, 115 filename=frame_ts, 116 ) 117 if not res["ok"]: 118 logging.error(f"Failed to post_match: {res}") 119 return 120 except Exception as ex: 121 logging.error(f"Slack files_upload_v2 failed with exception: {ex}") 122 123 file_id = res.get("files")[0].get("id") 124 attempt = 0 125 126 while attempt < 10: 127 attempt += 1 128 try: 129 res = await client.files_info(file=file_id) 130 if not res["ok"]: 131 logging.error(f"Failed to files_info: {res}") 132 return 133 thread_ts = res.get("file").get("shares").get("private").get(chan_id)[0].get("ts") 134 return thread_ts 135 except Exception as ex: 136 logging.error(f"failed to get files_info for attempt {attempt}") 137 await asyncio.sleep(3) 138 139 140 async def main_task(args): 141 archive = Archive(args.db_file) 142 client = AsyncWebClient(token=os.environ["SLACK_API_TOKEN"]) 143 sampler = StreamSampler("mavericksov", args.data_dir) 144 145 frames_q = asyncio.Queue() # frames that should be inspected. 146 archive_q = asyncio.Queue() # matches that should be archived. 147 announce_q = asyncio.Queue() # matches that should be announced in Slack. 148 149 async with asyncio.TaskGroup() as tg: 150 tg.create_task(sample_stream(frames_q, sampler, args.min_delay, args.max_delay, args.day_start, args.day_end)) 151 tg.create_task(find_matches(frames_q, archive_q, announce_q)) 152 tg.create_task(archive_matches(archive_q, archive)) 153 tg.create_task(announce_matches(announce_q, client)) 154 155 156 def main(): 157 parser = argparse.ArgumentParser() 158 parser.add_argument("-d", "--data-dir", default="./data", type=str) 159 parser.add_argument("-f", "--db-file", default="./archive.db", type=str) 160 parser.add_argument("-v", "--verbose", action="store_true") 161 parser.add_argument("--min-delay", default=5, type=int) 162 parser.add_argument("--max-delay", default=30, type=int) 163 parser.add_argument("--day-start", default="6:00", type=lambda s: datetime.strptime(s, "%H:%M").time()) 164 parser.add_argument("--day-end", default="21:00", type=lambda s: datetime.strptime(s, "%H:%M").time()) 165 args = parser.parse_args() 166 167 level = logging.INFO if args.verbose else logging.WARN 168 logging.basicConfig(level=level, format="%(asctime)s %(message)s") 169 170 asyncio.run(main_task(args))