SHARE
TWEET

Untitled

a guest Mar 22nd, 2019 63 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python3
  2.  
  3. import base64
  4. import json
  5. import os
  6. import io
  7. import datetime
  8. import time
  9. import cv2
  10. import boto3
  11. import botocore.response
  12.  
  13. threshold = 0.10
  14.  
  15. object_categories = [
  16.     'tug',
  17.     'bridge',
  18.     'aircraft',
  19.     'beltloader',
  20.     'fueltruck',
  21.     'towbar',
  22.     'towtug',
  23.     'truck',
  24.     'liftloader',
  25.     'gpu'
  26. ]
  27.  
  28.  
  29. def video_to_frames(video, path_output_dir):
  30.     capture = cv2.VideoCapture(video)
  31.     frames = 0
  32.     while True:
  33.         success, image = capture.read()
  34.         if success:
  35.             cv2.imwrite(os.path.join(path_output_dir,
  36.                                      'frame-%04d.jpg') % frames, image)
  37.             frames += 1
  38.         else:
  39.             break
  40.     capture.release()
  41.  
  42.  
  43. def read_file(file):
  44.     with open(file, "rb") as input:
  45.         return input.read()
  46.  
  47.  
  48. def save(filepath, fig=None):
  49.     '''Save the current image with no whitespace
  50.     Example filepath: "myfig.png" or r"C:\myfig.pdf"
  51.     '''
  52.     import matplotlib.pyplot as plt
  53.     if not fig:
  54.         fig = plt.gcf()
  55.  
  56.     plt.subplots_adjust(0,0,1,1,0,0)
  57.     for ax in fig.axes:
  58.         ax.axis('off')
  59.         ax.margins(0,0)
  60.         ax.xaxis.set_major_locator(plt.NullLocator())
  61.         ax.yaxis.set_major_locator(plt.NullLocator())
  62.  
  63.     basedir  = os.path.dirname(filepath)
  64.     basefile = os.path.basename(filepath)
  65.     basename = os.path.splitext(basefile)[0]
  66.     baseext = os.path.splitext(basefile)[1]
  67.  
  68.     fig.savefig(os.path.join(basedir, f"{basename}-detections{baseext}"), pad_inches = 0, bbox_inches='tight')
  69.  
  70. def visualize_detection(img_file, dets, classes=[], thresh=0.6):
  71.     import random
  72.     import matplotlib.pyplot as plt
  73.     import matplotlib.image as mpimg
  74.     img = mpimg.imread(img_file)
  75.     IMAGE_SIZE = (24, 16)
  76.     plt.figure(figsize=IMAGE_SIZE, frameon=False)
  77.     plt.imshow(img)
  78.     height = img.shape[0]
  79.     width = img.shape[1]
  80.     colors = dict()
  81.     for det in dets:
  82.         (klass, score, x0, y0, x1, y1) = det
  83.         if score < thresh:
  84.             continue
  85.         cls_id = int(klass)
  86.         if cls_id not in colors:
  87.             colors[cls_id] = (0, 1, 0)
  88.         xmin = int(x0 * width)
  89.         ymin = int(y0 * height)
  90.         xmax = int(x1 * width)
  91.         ymax = int(y1 * height)
  92.         rect = plt.Rectangle((xmin, ymin), xmax - xmin,
  93.                              ymax - ymin, fill=False,
  94.                              edgecolor=colors[cls_id],
  95.                              linewidth=1.5)
  96.         plt.gca().add_patch(rect)
  97.         class_name = str(cls_id)
  98.         if classes and len(classes) > cls_id:
  99.             class_name = classes[cls_id]
  100.         plt.gca().text(xmin, ymin - 2,
  101.                        '{:s} {:.3f}'.format(class_name, score),
  102.                        bbox=dict(facecolor=colors[cls_id], alpha=0.5),
  103.                        fontsize=12, color='white')
  104.     save(img_file, plt.gcf())
  105.     plt.close('all')
  106.  
  107. def main():
  108.     # Video to frames
  109.     # source = 'resources/target.mp4'
  110.     # destination = 'out'
  111.     # os.makedirs(destination, exist_ok=True)
  112.     # video_to_frames(source, destination)
  113.  
  114.     # Frames to predictions
  115.     sagemaker = boto3.client('sagemaker-runtime')
  116.  
  117.     for subdir, dirs, files in os.walk('out'):
  118.         for file in sorted(files):
  119.             if file.startswith(("frame")) and file.endswith((".jpg")):
  120.                 frame = os.path.join(subdir, file)
  121.  
  122.                 print(frame)
  123.  
  124.                 response = sagemaker.invoke_endpoint(
  125.                     EndpointName='f360-cavu-prod',
  126.                     Body=read_file(frame),
  127.                     ContentType='application/x-image',
  128.                     Accept='application/json',
  129.                 )
  130.                 body = response["Body"]
  131.                 body = body.read().decode('utf-8')
  132.                 body = json.loads(body)
  133.                 visualize_detection(frame, body['prediction'], object_categories, 0.10)
  134.  
  135.  
  136. if __name__ == '__main__':
  137.     main()
  138.     exit(0)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top